Compare commits

...
Author SHA1 Message Date
Emil 049113cec8 Harden the surface: 0600 databases, checksum-verified installs, rate-limited login and key exchange
coordinator / test (push) Waiting to run
python / test (push) Waiting to run
release / binaries (amd64, darwin) (push) Waiting to run
release / binaries (amd64, linux) (push) Waiting to run
release / binaries (amd64, windows) (push) Waiting to run
release / binaries (arm64, darwin) (push) Waiting to run
release / binaries (arm64, linux) (push) Waiting to run
release / binaries (arm64, windows) (push) Waiting to run
release / wheel (push) Waiting to run
release / release (push) Blocked by required conditions
release / image (push) Waiting to run
users / test (push) Waiting to run
2026-08-03 18:40:27 +03:00
Emil 330f95a375 Drop a leftover helper
coordinator / test (push) Waiting to run
python / test (push) Waiting to run
release / binaries (amd64, darwin) (push) Waiting to run
release / binaries (amd64, linux) (push) Waiting to run
release / binaries (amd64, windows) (push) Waiting to run
release / binaries (arm64, darwin) (push) Waiting to run
release / binaries (arm64, linux) (push) Waiting to run
release / binaries (arm64, windows) (push) Waiting to run
release / wheel (push) Waiting to run
release / release (push) Blocked by required conditions
release / image (push) Waiting to run
users / test (push) Waiting to run
2026-08-03 18:18:19 +03:00
Emil e744e62d03 Fix network URLs in serve: loopback agents, own-origin public URLs, longer transfer timeouts 2026-08-03 18:17:56 +03:00
Emil b9f47e556c Add uninstallers for the coordinator and the worker (release assets)
coordinator / test (push) Waiting to run
python / test (push) Waiting to run
release / binaries (amd64, darwin) (push) Waiting to run
release / binaries (amd64, linux) (push) Waiting to run
release / binaries (amd64, windows) (push) Waiting to run
release / binaries (arm64, darwin) (push) Waiting to run
release / binaries (arm64, linux) (push) Waiting to run
release / binaries (arm64, windows) (push) Waiting to run
release / wheel (push) Waiting to run
release / release (push) Blocked by required conditions
release / image (push) Waiting to run
users / test (push) Waiting to run
2026-08-03 18:04:04 +03:00
Emil e6b489f117 Log version control and isolation work in the session goal 2026-08-03 16:45:50 +03:00
Emil 73dde99e3a Run the wizard checks and task runners in isolated python mode (-I)
coordinator / test (push) Waiting to run
python / test (push) Waiting to run
release / binaries (amd64, darwin) (push) Waiting to run
release / binaries (amd64, linux) (push) Waiting to run
release / binaries (amd64, windows) (push) Waiting to run
release / binaries (arm64, darwin) (push) Waiting to run
release / binaries (arm64, linux) (push) Waiting to run
release / binaries (arm64, windows) (push) Waiting to run
release / wheel (push) Waiting to run
release / release (push) Blocked by required conditions
release / image (push) Waiting to run
users / test (push) Waiting to run
2026-08-03 16:42:38 +03:00
Emil ecc8944006 Drop an unused test helper
coordinator / test (push) Waiting to run
python / test (push) Waiting to run
release / binaries (amd64, darwin) (push) Waiting to run
release / binaries (amd64, linux) (push) Waiting to run
release / binaries (amd64, windows) (push) Waiting to run
release / binaries (arm64, darwin) (push) Waiting to run
release / binaries (arm64, linux) (push) Waiting to run
release / binaries (arm64, windows) (push) Waiting to run
release / wheel (push) Waiting to run
release / release (push) Blocked by required conditions
release / image (push) Waiting to run
users / test (push) Waiting to run
2026-08-03 16:36:00 +03:00
Emil dc15e2d04b Warn in the wizard preflight when the installed scimesh version mismatches the binary 2026-08-03 16:35:37 +03:00
Emil ff1fc25d77 Log worker concurrency in the session goal 2026-08-03 15:45:11 +03:00
Emil 6b67326c3b Run claim-execute-upload loops concurrently in the worker agent
coordinator / test (push) Waiting to run
python / test (push) Waiting to run
release / binaries (amd64, darwin) (push) Waiting to run
release / binaries (amd64, linux) (push) Waiting to run
release / binaries (amd64, windows) (push) Waiting to run
release / binaries (arm64, darwin) (push) Waiting to run
release / binaries (arm64, linux) (push) Waiting to run
release / binaries (arm64, windows) (push) Waiting to run
release / wheel (push) Waiting to run
release / release (push) Blocked by required conditions
release / image (push) Waiting to run
users / test (push) Waiting to run
2026-08-03 15:36:22 +03:00
Emil d57f8778ac Log the parallel search workload in the session goal 2026-08-03 15:20:10 +03:00
Emil 5b2d5b5f6e Add similarity-search-parallel: thread-pool scoring, byte-identical to similarity-search
coordinator / test (push) Waiting to run
python / test (push) Waiting to run
release / binaries (amd64, darwin) (push) Waiting to run
release / binaries (amd64, linux) (push) Waiting to run
release / binaries (amd64, windows) (push) Waiting to run
release / binaries (arm64, darwin) (push) Waiting to run
release / binaries (arm64, linux) (push) Waiting to run
release / binaries (arm64, windows) (push) Waiting to run
release / wheel (push) Waiting to run
release / release (push) Blocked by required conditions
release / image (push) Waiting to run
users / test (push) Waiting to run
2026-08-03 15:12:18 +03:00
Emil b63441bb8c Revert "Warn about non-ASCII characters in text fields of the job form"
coordinator / test (push) Waiting to run
python / test (push) Waiting to run
release / binaries (amd64, darwin) (push) Waiting to run
release / binaries (amd64, linux) (push) Waiting to run
release / binaries (amd64, windows) (push) Waiting to run
release / binaries (arm64, darwin) (push) Waiting to run
release / binaries (arm64, linux) (push) Waiting to run
release / binaries (arm64, windows) (push) Waiting to run
release / wheel (push) Waiting to run
release / release (push) Blocked by required conditions
release / image (push) Waiting to run
users / test (push) Waiting to run
This reverts commit 91d049b4df.
2026-08-03 14:34:46 +03:00
Emil 91d049b4df Warn about non-ASCII characters in text fields of the job form 2026-08-03 14:34:12 +03:00
37 changed files with 1644 additions and 40 deletions
+2
View File
@@ -110,6 +110,8 @@ jobs:
artifacts/*
install.sh
install.ps1
uninstall.sh
uninstall.ps1
# Pre-release tags (e.g. v1.1.0-alpha.1) publish as pre-releases.
prerelease: ${{ contains(github.ref_name, '-alpha') }}
generate_release_notes: true
+26
View File
@@ -38,6 +38,22 @@
## Plan (предыдущая задача — выполнена)
17. Docker E2E пайплайна «install как человек → serve → визард → воркер → джоб» — выполнено, см. Progress ниже.
## Progress (день: конкурентность в агенте — параллелизм через нашу архитектуру)
- ✅ По просьбе «реализовать параллелизм через нашу архитектуру»: `WORKER_CONCURRENCY` (поле `concurrency` в конфиге визарда, шаг Machine) — агент регистрируется один раз и ведёт N циклов claim→выполнение→upload под одним worker id. N шардов обрабатываются параллельно на одной машине, используя таски координатора как единицу параллелизма; SDK не менялся.
- ✅ Реализация: `Config.Concurrency` + env; `Daemon.RunForever` → register once → N goroutine-циклов (общий счётчик MaxTasks под мьютексом); визард: поле «Concurrent task loops» + конфиг; тест с маркер-скриптом доказывает 3 параллельных исполнения (race-тесты зелёные после мьютекса в fake).
- ✅ Измерено на релизном коде в Docker (один воркер, 8 шардов): concurrency=1 → 12s; concurrency=4 → 4s (**3×**). Плюс прежний `similarity-search-parallel` (потоки внутри шарда) композируется с конкурентностью.
- ✅ Релиз v1.1.0-alpha.20 (бинарники + wheel); полный гейт: race + lint 0 issues + pytest 213.
- ✅ Бинарник worker-agent на машине пользователя обновлён до alpha.20.
- GIL-высвобождения в RDKit нет (проверено: `RDKIT_ALLOW_GIL_RELEASE` не помогает), поэтому внутришардовые потоки не ускоряют чистый RDKit-путь — конкурентность задач это и компенсирует.
## Progress (утро: similarity-search-parallel)
- ✅ Новый workload `similarity-search-parallel@1.0.0` (отдельная версия, как просил пользователь): подкласс `SimilaritySearchSDKWorkload` + параллельное ядро `search_parallel/core.py` — fingerprinting+скоринг шарда через `ThreadPoolExecutor` (параметр `threads`, default CPU count); `pool.map` сохраняет порядок строк, поэтому merge идентичен последовательному (`_HeapEntry`) и результат **байт-в-байт** равен `similarity-search` при любом числе потоков.
- ✅ Тесты: байт-в-байт vs эталон для threads 1/2/4 с намеренными связями (изомеры, дубликаты), executor-прогон, валидация параметров, регистрация манифеста; 213 pytest зелёные.
- ✅ Экспортирован в каталог координатора (5 ворклоадов), Go-тесты зелёные; release v1.1.0-alpha.19 (бинарники + wheel `scimesh-1.1.0a19`).
- ✅ Распределённый E2E в Docker: воркер с wheel a19, джоб similarity-search-parallel (threads=4) → completed → результат байт-в-байт = локальному эталону.
- ✅ Бинарники на машине пользователя обновлены до alpha.19 (для появления ворклоада в UI нужен рестарт serve + переустановка рантайма воркера).
- Примечание: в CPython потоки не ускоряют чистый RDKit-путь (GIL), но структура готова к ядрам, отпускающим GIL (numpy и т.п.); при желании можно добавить процесс-пул как отдельный workload в будущей версии протокола.
## Progress (ночная сессия)
-**П.1 Пустое имя воркера**: `domain.NewWorker` нормализует/отклоняет пустое имя + `TestNewWorkerRejectsBlankName` (починен `fixedTime``testNow`).
-**П.2 `--check`**: пробует managed venv (если установлен) + реальный пробинг учётки (exchange ключа / claim-пробa) — `CheckAuth` + тесты; на машине пользователя: `✓ auth: credential accepted`, venv python, scimesh installed.
@@ -59,8 +75,18 @@
- ✅ На машине пользователя: визард alpha.15, правильный токен, venv из wheel, воркер emil-pc online, 15 пустых воркеров вычищены из БД.
- ✅ Гейт: race + lint + pytest 208.
## Progress (день 2: контроль версий + изоляция)
- ✅ Визард сверяет установленную версию scimesh с версией бинарника: preflight показывает «installed X, worker needs Y — press Install» и снова предлагает кнопку Install (ensureMatchingScimeshVersion, dev-сборки пропускают). Проверено на машине: a21 vs a22 → красный → Install → зелёный 1.1.0a22.
- ✅ Все проверки и task-раннеры запускаются с `-I` (isolated): каталог запуска больше не перекрывает venv (был артефакт cwd-затенения исходниками). Затронуты: check, дефолтные раннеры агента, визард (config/save/start), serve, agent_cmd.
- ✅ Релиз v1.1.0-alpha.22; полный гейт (race + lint 0 issues + pytest 213); бинарник юзера обновлён, рантайм переустановлен (a22), воркер с concurrency=4; большой джоб юзера (1934 шарда) завершён (очередь 0).
## Completion
COMPLETED — ночной план выполнен полностью (10 пунктов + 2 найденных бага, включая E2E quorum на релизном коде). Все гейты зелёные, релиз v1.1.0-alpha.16 опубликован.
## Progress (день 2: контроль версий + изоляция)
- ✅ Визард сверяет установленную версию scimesh с версией бинарника: preflight показывает «installed X, worker needs Y — press Install» и снова предлагает кнопку Install (ensureMatchingScimeshVersion, dev-сборки пропускают). Проверено на машине: a21 vs a22 → красный → Install → зелёный 1.1.0a22.
- ✅ Все проверки и task-раннеры запускаются с `-I` (isolated): каталог запуска больше не перекрывает venv (был артефакт cwd-затенения исходниками). Затронуты: check, дефолтные раннеры агента, визард (config/save/start), serve, agent_cmd.
- ✅ Релиз v1.1.0-alpha.22; полный гейт (race + lint 0 issues + pytest 213); бинарник юзера обновлён, рантайм переустановлен (a22), воркер с concurrency=4; большой джоб юзера (1934 шарда) завершён (очередь 0).
## Completion (предыдущая задача)
COMPLETED — пайплайн доведён до рабочего состояния и проверен на релизных артефактах v1.1.0-alpha.14.
+10 -1
View File
@@ -61,7 +61,16 @@ powershell -ExecutionPolicy Bypass -c "irm https://raw.githubusercontent.com/emi
```
Set `SCIMESH_AUTO_START=0` to install without starting anything. The old demo
control room was removed: `/ui` is the admin console. A standalone
control room was removed: `/ui` is the admin console.
To remove a component, run the matching uninstaller (data is kept unless you
pass `--purge`):
```bash
curl -fsSL https://raw.githubusercontent.com/emil28092005/SciMesh/main/uninstall.sh | bash -s coordinator
curl -fsSL https://raw.githubusercontent.com/emil28092005/SciMesh/main/uninstall.sh | bash -s worker --purge
# Windows: irm .../uninstall.ps1 | iex (set $env:SCIMESH_COMPONENT, -Purge deletes data)
``` A standalone
worker is installed the same way (`bash -s worker`, or
`SCIMESH_COMPONENT=worker` on Windows); its installer opens the local setup
wizard (`worker-agent setup`) in the browser automatically.
+1 -1
View File
@@ -45,7 +45,7 @@ func runAgent(args []string) error {
return fmt.Errorf("--coordinator-url, --token, and --work-dir are required")
}
if *taskRunner == "" {
*taskRunner = "python -m scimesh.worker.task"
*taskRunner = "python -I -m scimesh.worker.task"
}
logger := slog.New(slog.NewTextHandler(os.Stderr, nil))
+44 -12
View File
@@ -10,6 +10,7 @@ import (
"flag"
"fmt"
"log/slog"
"net"
"os"
"os/exec"
"path/filepath"
@@ -96,18 +97,16 @@ func runServe(args []string) error {
defer func() { _ = closeUsers() }()
// 5. Local worker agents before the server, so they can claim immediately.
coordinatorURL := "http://" + *addr
agents, err := spawnAgents(ctx, log, *dataDir, *workers, coordinatorURL, workerToken, venvPython)
// They always dial the loopback address: an --addr of 0.0.0.0 is not a
// connectable target from the same host.
agentURL, resolvedPublic := serveURLs(*addr, *publicURL)
agents, err := spawnAgents(ctx, log, *dataDir, *workers, agentURL, workerToken, venvPython)
if err != nil {
return err
}
defer stopAgents(agents)
// 6. The coordinator server itself.
coordinatorPublicURL := *publicURL
if coordinatorPublicURL == "" {
coordinatorPublicURL = "http://" + *addr
}
cfg := infra.Config{
Addr: *addr,
DatabaseEngine: "sqlite",
@@ -115,8 +114,10 @@ func runServe(args []string) error {
Token: workerToken,
JWTSecret: jwtSecret,
UserserviceURL: "http://" + usersAddr,
PublicCoordinatorURL: coordinatorPublicURL,
PublicUserserviceURL: "http://" + usersAddr,
PublicCoordinatorURL: resolvedPublic,
// The exchange is fronted by the coordinator's own proxy, so the UI
// falls back to the coordinator origin for USERSERVICE_URL.
PublicUserserviceURL: "",
LogLevel: "info",
StorageDir: filepath.Join(*dataDir, "artifacts"),
DocsDir: *docsDir,
@@ -132,12 +133,13 @@ func runServe(args []string) error {
WorkerOfflineAfter: 1 * time.Minute,
AutoMigrate: true,
}
browserURL, _ := serveURLs(*addr, *publicURL)
if *open {
openBrowser("http://" + *addr + "/ui/admin")
openBrowser(browserURL + "/ui/admin")
}
// Print the login once the server is about to start.
fmt.Printf("\nSciMesh is starting at http://%s/ui\n", *addr)
fmt.Printf("\nSciMesh is starting at %s/ui\n", browserURL)
fmt.Printf(" admin login: %s / %s\n", *email, *password)
if runtimeStatus(venvPython) {
fmt.Printf(" scientific runtime: ready (%s)\n", venvPython)
@@ -236,9 +238,9 @@ func stopAgents(agents []*exec.Cmd) {
// system `python`.
func defaultTaskRunner(venvPython string) string {
if runtimeStatus(venvPython) {
return venvPython + " -m scimesh.worker.task"
return venvPython + " -I -m scimesh.worker.task"
}
return "python -m scimesh.worker.task"
return "python -I -m scimesh.worker.task"
}
// ensureRuntime creates the managed venv and installs scimesh into it, unless
@@ -337,3 +339,33 @@ func openBrowser(target string) {
// #nosec G204 -- target is the local UI URL the operator asked to open.
_ = exec.CommandContext(context.Background(), command, target).Start()
}
// serveURLs derives the two addresses of a serve instance from the listen
// address and the optional --public-url flag:
//
// - the agent URL is always the loopback form of the port, because spawned
// local workers share the host and 0.0.0.0 is not connectable from it;
// - the public URL is what browsers and remote workers are told. An explicit
// --public-url wins; a listen host that is a real address is used as-is;
// a wildcard host (0.0.0.0, ::, or empty) yields an empty public URL, so
// the UI falls back to the browser's own origin (the coordinator's LAN
// address as the browser sees it).
func serveURLs(addr, publicURL string) (agentURL, resolvedPublic string) {
host, port, err := net.SplitHostPort(addr)
if err != nil {
// No port in the listen address: assume the default and treat the
// whole string as a host (e.g. a bare wildcard).
host, port = addr, "8080"
}
agentURL = "http://127.0.0.1:" + port
if publicURL != "" {
return agentURL, publicURL
}
host = strings.Trim(host, "[]")
switch host {
case "", "0.0.0.0", "::":
return agentURL, ""
default:
return agentURL, "http://" + addr
}
}
@@ -0,0 +1,22 @@
package main
import "testing"
func TestServeURLs(t *testing.T) {
cases := []struct {
addr, public, agent, resolved string
}{
{"127.0.0.1:8080", "", "http://127.0.0.1:8080", "http://127.0.0.1:8080"},
{"0.0.0.0:8080", "", "http://127.0.0.1:8080", ""},
{":8080", "", "http://127.0.0.1:8080", ""},
{"::", "", "http://127.0.0.1:8080", ""},
{"192.168.1.10:8080", "", "http://127.0.0.1:8080", "http://192.168.1.10:8080"},
{"0.0.0.0:8080", "http://cluster.example:8080", "http://127.0.0.1:8080", "http://cluster.example:8080"},
}
for _, c := range cases {
agent, resolved := serveURLs(c.addr, c.public)
if agent != c.agent || resolved != c.resolved {
t.Errorf("serveURLs(%q, %q) = (%q, %q), want (%q, %q)", c.addr, c.public, agent, resolved, c.agent, c.resolved)
}
}
}
+5 -1
View File
@@ -88,8 +88,12 @@ func CheckEnvironment(ctx context.Context) CheckReport {
// execute with.
func CheckEnvironmentWithPython(ctx context.Context, python string) CheckReport {
report := CheckReport{Agent: Version, Python: CheckItem{Name: "python", OK: true, Detail: python}}
// The version comes from importlib.metadata, so the wizard can compare the
// installed package with the binary version and offer an upgrade. -I keeps
// the working directory out of sys.path, so a scimesh checkout in the
// wizard's cwd can never shadow the venv installation.
//nolint:gosec // G204: python is a resolved interpreter path, the argument list is constant
cmd := exec.CommandContext(ctx, python, "-c", "import scimesh; print(scimesh.__version__ if hasattr(scimesh, '__version__') else 'installed')")
cmd := exec.CommandContext(ctx, python, "-I", "-c", "import importlib.metadata as m; print(m.version('scimesh'))")
out, err := cmd.Output()
if err != nil {
// The worker executes workloads by spawning scimesh's task runner, so
+7 -1
View File
@@ -48,6 +48,12 @@ type Client struct {
}
func NewClient(baseURL string, tokens TokenProvider, timeout time.Duration) *Client {
// Payload transfers get a more generous budget than control calls: a large
// shard over a slow link easily outlives the API timeout.
transferTimeout := timeout * 4
if transferTimeout < 2*time.Minute {
transferTimeout = 2 * time.Minute
}
return &Client{
baseURL: strings.TrimRight(baseURL, "/"),
tokens: tokens,
@@ -57,7 +63,7 @@ func NewClient(baseURL string, tokens TokenProvider, timeout time.Duration) *Cli
CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse },
},
dlClient: &http.Client{
Timeout: timeout,
Timeout: transferTimeout,
CheckRedirect: func(req *http.Request, via []*http.Request) error {
if len(via) >= 10 {
return fmt.Errorf("too many redirects")
+14
View File
@@ -223,3 +223,17 @@ func sha256Of(t *testing.T, value string) string {
digest := sha256.Sum256([]byte(value))
return fmt.Sprintf("%x", digest)
}
func TestNewClientTransferTimeoutExceedsAPITimeout(t *testing.T) {
c := NewClient("http://coord:8080", &StaticToken{token: "t"}, 30*time.Second)
if c.apiClient.Timeout != 30*time.Second {
t.Errorf("api timeout = %v, want 30s", c.apiClient.Timeout)
}
if c.dlClient.Timeout < 2*time.Minute {
t.Errorf("transfer timeout = %v, want at least 2m", c.dlClient.Timeout)
}
short := NewClient("http://coord:8080", &StaticToken{token: "t"}, 3*time.Minute)
if short.dlClient.Timeout != 12*time.Minute {
t.Errorf("transfer timeout = %v, want 4x the api timeout", short.dlClient.Timeout)
}
}
+18 -1
View File
@@ -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) {
@@ -108,7 +112,7 @@ func LoadConfig() (*Config, error) {
return nil, err
}
if len(runner) == 0 {
runner = []string{"python", "-m", "scimesh.worker.task"}
runner = []string{"python", "-I", "-m", "scimesh.worker.task"}
}
maxTasks := 0
if raw := os.Getenv("MAX_TASKS"); raw != "" {
@@ -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 == "" {
+6 -1
View File
@@ -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,11 +104,15 @@ 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
}
if len(config.TaskRunner) == 0 {
config.TaskRunner = []string{"python", "-m", "scimesh.worker.task"}
config.TaskRunner = []string{"python", "-I", "-m", "scimesh.worker.task"}
}
config.PollInterval = 2 * time.Second
config.RequestTimeout = 30 * time.Second
+47 -7
View File
@@ -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
}
+79
View File
@@ -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)
}
}
+33 -2
View File
@@ -323,7 +323,7 @@ func (s *Server) ensureVenvTaskRunner() {
return
}
if venv := s.venvPython(); venv != "" {
file.TaskRunner = []string{venv, "-m", "scimesh.worker.task"}
file.TaskRunner = []string{venv, "-I", "-m", "scimesh.worker.task"}
if payload, err := json.MarshalIndent(file, "", " "); err == nil {
_ = os.WriteFile(s.cfgPath, append(payload, '\n'), 0o600)
}
@@ -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,12 +420,15 @@ 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.
if len(file.TaskRunner) == 0 {
if venv := s.venvPython(); venv != "" {
file.TaskRunner = []string{venv, "-m", "scimesh.worker.task"}
file.TaskRunner = []string{venv, "-I", "-m", "scimesh.worker.task"}
}
}
if err := agent.SaveConfigFile(s.cfgPath, file); err != nil {
@@ -449,6 +454,9 @@ func (s *Server) handleTest(w http.ResponseWriter, r *http.Request) {
// checking the bare system python3 would keep reporting scimesh as
// missing even though the worker would run with the venv.
report := agent.RunCheck(r.Context(), url, s.venvPython(), req.Token, req.WorkerKey, req.UserserviceURL)
if report.Scimesh.OK {
report.Scimesh = ensureMatchingScimeshVersion(report.Scimesh)
}
writeJSON(w, http.StatusOK, report)
}
@@ -623,3 +631,26 @@ func truncate(s string, n int) string {
}
return s[:n] + "…"
}
// ensureMatchingScimeshVersion flips a green scimesh check to a stale one when
// the installed package does not match the worker-agent's own version: a
// version-locked wheel is the only supported runtime, and a mismatch means the
// workload catalog the worker advertises is not what it executes. The wizard
// UI then offers the Install button again. Dev builds have no release wheel,
// so they skip the comparison.
func ensureMatchingScimeshVersion(item agent.CheckItem) agent.CheckItem {
if agent.Version == "" || agent.Version == "dev" {
return item
}
want := agent.NormalizePEP440(agent.Version)
got := strings.TrimSpace(item.Detail)
if got == "" || got == want {
return item
}
item.OK = false
item.Detail = fmt.Sprintf(
"installed scimesh %s, but this worker-agent (%s) needs %s — press Install to upgrade",
got, agent.Version, want,
)
return item
}
@@ -429,8 +429,8 @@ func TestStartPinsTheVenvTaskRunner(t *testing.T) {
if err != nil {
t.Fatal(err)
}
if len(config.TaskRunner) != 3 || config.TaskRunner[0] != venvPython || config.TaskRunner[1] != "-m" || config.TaskRunner[2] != "scimesh.worker.task" {
t.Errorf("task runner = %v, want the venv python runner", config.TaskRunner)
if len(config.TaskRunner) != 4 || config.TaskRunner[0] != venvPython || config.TaskRunner[1] != "-I" || config.TaskRunner[2] != "-m" || config.TaskRunner[3] != "scimesh.worker.task" {
t.Errorf("task runner = %v, want the venv python runner with -I", config.TaskRunner)
}
}
@@ -450,8 +450,8 @@ func TestSaveConfigPinsVenvRunnerWhenPresent(t *testing.T) {
if err != nil {
t.Fatal(err)
}
if len(config.TaskRunner) != 3 || config.TaskRunner[0] != venvPython {
t.Errorf("task runner = %v, want the venv python", config.TaskRunner)
if len(config.TaskRunner) != 4 || config.TaskRunner[0] != venvPython || config.TaskRunner[1] != "-I" {
t.Errorf("task runner = %v, want the venv python with -I", config.TaskRunner)
}
}
@@ -462,7 +462,7 @@ func TestTestProbesTheVenvPythonAfterInstall(t *testing.T) {
// a fake scimesh version so the preflight goes green through the venv.
venvPython := filepath.Join(server.dir, "venv", "bin", "python")
_ = os.MkdirAll(filepath.Dir(venvPython), 0o755)
_ = os.WriteFile(venvPython, []byte("#!/bin/sh\nif [ \"$1\" = \"-c\" ]; then echo 9.9.9-test; exit 0; fi\nexit 0\n"), 0o755)
_ = os.WriteFile(venvPython, []byte("#!/bin/sh\nfor a in \"$@\"; do if [ \"$a\" = \"-c\" ]; then echo 9.9.9-test; exit 0; fi; done\nexit 0\n"), 0o755)
req, _ := http.NewRequestWithContext(context.Background(), http.MethodPost, base+"/api/test", strings.NewReader(`{"coordinator_url":"http://127.0.0.1:1"}`))
req.Header.Set("Content-Type", "application/json")
@@ -497,3 +497,26 @@ time=6 level=WARN msg="agent cycle failed" error="boom"
t.Errorf("stats = %+v, want registered claimed=2 completed=1 failed=1", stats)
}
}
func testCheckScimeshVersion(t *testing.T, installed, binary string, wantOK bool, wantDetail string) {
t.Helper()
old := agent.Version
agent.Version = binary
t.Cleanup(func() { agent.Version = old })
item := ensureMatchingScimeshVersion(agent.CheckItem{Name: "scimesh", OK: true, Detail: installed})
if item.OK != wantOK {
t.Errorf("installed=%s binary=%s: ok=%v, want %v (%s)", installed, binary, item.OK, wantOK, item.Detail)
}
if wantDetail != "" && !strings.Contains(item.Detail, wantDetail) {
t.Errorf("detail = %q, want it to contain %q", item.Detail, wantDetail)
}
}
func TestEnsureMatchingScimeshVersion(t *testing.T) {
testCheckScimeshVersion(t, "1.1.0a20", "1.1.0-alpha.20", true, "")
testCheckScimeshVersion(t, "1.1.0a17", "1.1.0-alpha.20", false, "press Install to upgrade")
testCheckScimeshVersion(t, "1.1.0a16.dev7+gea0fb8c59.d20260803", "1.1.0-alpha.20", false, "needs 1.1.0a20")
// Dev builds and unknown versions never block.
testCheckScimeshVersion(t, "anything", "dev", true, "")
testCheckScimeshVersion(t, "1.1.0a20", "", true, "")
}
@@ -141,6 +141,7 @@ code{font-family:var(--mono);font-size:.86em}
</div>
</div>
<div class="field" id="cpu-field" style="display:none"><label>CPU count</label><input id="in-cpu" type="number" min="1" value="1"></div>
<div class="field"><label>Concurrent task loops</label><input id="in-conc" type="number" min="1" max="64" value="1"><p class="hint">Process this many shards in parallel on this machine. Each loop runs its own task runner subprocess.</p></div>
<div class="actions"><button class="btn btn-ghost" id="b2b">← Back</button><button class="btn btn-primary" id="b2">Continue →</button></div>
</div>
</div>
@@ -228,9 +229,10 @@ 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'];
if(state.venvPython)cfg.task_runner=[state.venvPython,'-I','-m','scimesh.worker.task'];
return cfg;
}
@@ -18,6 +18,7 @@ import (
"encoding/json"
"errors"
"fmt"
"os"
"time"
_ "modernc.org/sqlite"
@@ -37,9 +38,29 @@ func Open(path string) (*sql.DB, error) {
_ = db.Close()
return nil, fmt.Errorf("ping sqlite database: %w", err)
}
if err := lockDownDatabase(path); err != nil {
_ = db.Close()
return nil, fmt.Errorf("lock down sqlite database: %w", err)
}
return db, nil
}
// lockDownDatabase restricts the database files to the owner: sqlite creates
// them with the process umask (0644), which would let any local user read job
// metadata and password hashes. WAL/SHM siblings inherit the main file's mode,
// so existing ones are corrected too. Best-effort: failures only warn callers
// via the returned error, never corrupt state.
func lockDownDatabase(path string) error {
for _, candidate := range []string{path, path + "-wal", path + "-shm"} {
if info, err := os.Stat(candidate); err == nil && !info.IsDir() {
if err := os.Chmod(candidate, 0o600); err != nil {
return err
}
}
}
return nil
}
// querier is satisfied by both *sql.DB and *sql.Tx, letting every repository
// method run identically inside or outside a transaction.
type querier interface {
@@ -4,6 +4,7 @@ import (
"context"
"database/sql"
"errors"
"os"
"path/filepath"
"testing"
"time"
@@ -347,3 +348,19 @@ func TestCancelByJobInvalidatesTasks(t *testing.T) {
t.Errorf("cancelled task = %+v", got)
}
}
func TestOpenRestrictsDatabasePermissions(t *testing.T) {
path := filepath.Join(t.TempDir(), "locked.db")
db, err := Open(path)
if err != nil {
t.Fatal(err)
}
_ = db.Close()
info, err := os.Stat(path)
if err != nil {
t.Fatal(err)
}
if perm := info.Mode().Perm(); perm != 0o600 {
t.Errorf("db perms = %o, want 600", perm)
}
}
@@ -0,0 +1,117 @@
package http
import (
"net"
"net/http"
"sync"
"time"
)
// documentedLimits describes the default policy for the two public surfaces;
// keep in sync with loginRatePerMinute and exchangeRatePerMinute below.
const (
loginRatePerMinute = 10
loginBurst = 5
exchangeRatePerMinute = 30
exchangeBurst = 10
)
// tokenBucket is a fixed-rate token bucket for one client address.
type tokenBucket struct {
mu sync.Mutex
tokens float64
last time.Time
rate float64 // tokens per second
burst float64
}
func newTokenBucket(ratePerMinute, burst float64) *tokenBucket {
return &tokenBucket{
tokens: burst,
last: time.Now(),
rate: ratePerMinute / 60,
burst: burst,
}
}
// allow consumes one token when available; the bucket refills continuously.
func (b *tokenBucket) allow() bool {
b.mu.Lock()
defer b.mu.Unlock()
now := time.Now()
b.tokens += now.Sub(b.last).Seconds() * b.rate
if b.tokens > b.burst {
b.tokens = b.burst
}
b.last = now
if b.tokens < 1 {
return false
}
b.tokens--
return true
}
// ipLimiter tracks one bucket per client address and prunes stale entries.
type ipLimiter struct {
mu sync.Mutex
buckets map[string]*tokenBucket
ratePerMinute float64
burst float64
}
func newIPLimiter(ratePerMinute, burst float64) *ipLimiter {
return &ipLimiter{
buckets: make(map[string]*tokenBucket),
ratePerMinute: ratePerMinute,
burst: burst,
}
}
// Allow reports whether the caller's address may proceed. It also sweeps
// entries idle for more than ten minutes so the map stays bounded.
func (l *ipLimiter) Allow(r *http.Request) bool {
ip := remoteIP(r)
l.mu.Lock()
if len(l.buckets) > 1024 {
cutoff := time.Now().Add(-10 * time.Minute)
for addr, bucket := range l.buckets {
bucket.mu.Lock()
idle := bucket.last.Before(cutoff)
bucket.mu.Unlock()
if idle {
delete(l.buckets, addr)
}
}
}
bucket, ok := l.buckets[ip]
if !ok {
bucket = newTokenBucket(l.ratePerMinute, l.burst)
l.buckets[ip] = bucket
}
l.mu.Unlock()
return bucket.allow()
}
func remoteIP(r *http.Request) string {
host, _, err := net.SplitHostPort(r.RemoteAddr)
if err != nil {
return r.RemoteAddr
}
return host
}
// rateLimited wraps a handler with a per-address limiter; exhausted callers
// receive 429 with a Retry-After header.
func rateLimited(limiter *ipLimiter, next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if !limiter.Allow(r) {
w.Header().Set("Retry-After", "60")
writeJSON(w, http.StatusTooManyRequests, map[string]string{
"error": "too many requests, try again shortly",
"request_id": requestIDFrom(r.Context()),
})
return
}
next.ServeHTTP(w, r)
})
}
@@ -0,0 +1,54 @@
package http
import (
"net/http"
"net/http/httptest"
"testing"
)
func TestTokenBucketBurstThenThrottles(t *testing.T) {
bucket := newTokenBucket(10, 3)
for i := 0; i < 3; i++ {
if !bucket.allow() {
t.Fatalf("request %d must pass within the burst", i)
}
}
if bucket.allow() {
t.Error("fourth request within the burst must be throttled")
}
}
func TestRateLimitedReturns429(t *testing.T) {
limiter := newIPLimiter(10, 2)
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusOK)
})
handler := rateLimited(limiter, next)
req := httptest.NewRequestWithContext(t.Context(), http.MethodPost, "/ui/login", nil)
req.RemoteAddr = "10.0.0.5:5555"
// Burst is 2: the first two pass, the third is throttled.
for i := 0; i < 2; i++ {
rec := httptest.NewRecorder()
handler.ServeHTTP(rec, req)
if rec.Code != http.StatusOK {
t.Fatalf("request %d: got %d", i, rec.Code)
}
}
rec := httptest.NewRecorder()
handler.ServeHTTP(rec, req)
if rec.Code != http.StatusTooManyRequests {
t.Fatalf("third request: got %d, want 429", rec.Code)
}
if rec.Header().Get("Retry-After") == "" {
t.Error("429 must carry Retry-After")
}
// A different address is not throttled by the same bucket.
other := httptest.NewRequestWithContext(t.Context(), http.MethodPost, "/ui/login", nil)
other.RemoteAddr = "10.0.0.6:5555"
rec = httptest.NewRecorder()
handler.ServeHTTP(rec, other)
if rec.Code != http.StatusOK {
t.Errorf("other client: got %d, want 200", rec.Code)
}
}
@@ -135,7 +135,7 @@ func (s *Server) Handler(token string, uiToken ...string) http.Handler {
// Worker-key exchange is fronted by the coordinator when the userservice
// is embedded (serve mode): the key itself is the credential.
if s.userserviceURL != "" {
mux.HandleFunc("POST /worker-tokens/exchange", s.handleWorkerTokenExchangeProxy)
mux.Handle("POST /worker-tokens/exchange", rateLimited(newIPLimiter(exchangeRatePerMinute, exchangeBurst), http.HandlerFunc(s.handleWorkerTokenExchangeProxy)))
}
hasBasicAuth := len(uiToken) > 0 && uiToken[0] != ""
@@ -163,7 +163,7 @@ func (s *Server) Handler(token string, uiToken ...string) http.Handler {
if s.uiSessionMode() {
// Public auth pages — reachable without a session so a user can log in.
ui.HandleFunc("GET /ui/login", s.handleUILoginForm)
ui.HandleFunc("POST /ui/login", s.handleUILogin)
ui.Handle("POST /ui/login", rateLimited(newIPLimiter(loginRatePerMinute, loginBurst), http.HandlerFunc(s.handleUILogin)))
ui.HandleFunc("GET /ui/logout-form", s.handleUILogoutForm)
ui.HandleFunc("GET /ui/register", s.handleUIRegisterForm)
ui.HandleFunc("POST /ui/register", s.handleUIRegister)
@@ -7,6 +7,7 @@ import (
"database/sql"
"errors"
"fmt"
"os"
"time"
_ "modernc.org/sqlite"
@@ -290,6 +291,17 @@ func indexOf(haystack, needle string) int {
}
// Open opens (and creates when missing) the userservice database file.
func lockDownDatabase(path string) error {
for _, candidate := range []string{path, path + "-wal", path + "-shm"} {
if info, err := os.Stat(candidate); err == nil && !info.IsDir() {
if err := os.Chmod(candidate, 0o600); err != nil {
return err
}
}
}
return nil
}
func Open(path string) (*sql.DB, error) {
dsn := fmt.Sprintf("file:%s?_pragma=busy_timeout(10000)&_pragma=journal_mode(WAL)&_pragma=foreign_keys(1)&_pragma=synchronous(NORMAL)", path)
db, err := sql.Open("sqlite", dsn)
@@ -300,5 +312,9 @@ func Open(path string) (*sql.DB, error) {
_ = db.Close()
return nil, fmt.Errorf("ping userservice database: %w", err)
}
if err := lockDownDatabase(path); err != nil {
_ = db.Close()
return nil, fmt.Errorf("lock down userservice database: %w", err)
}
return db, nil
}
@@ -0,0 +1,97 @@
package http
import (
"net"
"net/http"
"sync"
"time"
)
// The same per-address token-bucket policy as the coordinator transport:
// login is the credential brute-force surface, the exchange the only public
// token-minting one.
const (
loginRatePerMinute = 10
loginBurst = 5
exchangeRatePerMinute = 30
exchangeBurst = 10
)
type tokenBucket struct {
mu sync.Mutex
tokens float64
last time.Time
rate float64
burst float64
}
func newTokenBucket(ratePerMinute, burst float64) *tokenBucket {
return &tokenBucket{tokens: burst, last: time.Now(), rate: ratePerMinute / 60, burst: burst}
}
func (b *tokenBucket) allow() bool {
b.mu.Lock()
defer b.mu.Unlock()
now := time.Now()
b.tokens += now.Sub(b.last).Seconds() * b.rate
if b.tokens > b.burst {
b.tokens = b.burst
}
b.last = now
if b.tokens < 1 {
return false
}
b.tokens--
return true
}
type ipLimiter struct {
mu sync.Mutex
buckets map[string]*tokenBucket
ratePerMinute float64
burst float64
}
func newIPLimiter(ratePerMinute, burst float64) *ipLimiter {
return &ipLimiter{buckets: map[string]*tokenBucket{}, ratePerMinute: ratePerMinute, burst: burst}
}
func (l *ipLimiter) Allow(r *http.Request) bool {
host, _, err := net.SplitHostPort(r.RemoteAddr)
if err != nil {
host = r.RemoteAddr
}
l.mu.Lock()
if len(l.buckets) > 1024 {
cutoff := time.Now().Add(-10 * time.Minute)
for addr, bucket := range l.buckets {
bucket.mu.Lock()
idle := bucket.last.Before(cutoff)
bucket.mu.Unlock()
if idle {
delete(l.buckets, addr)
}
}
}
bucket, ok := l.buckets[host]
if !ok {
bucket = newTokenBucket(l.ratePerMinute, l.burst)
l.buckets[host] = bucket
}
l.mu.Unlock()
return bucket.allow()
}
func rateLimited(limiter *ipLimiter, next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if !limiter.Allow(r) {
w.Header().Set("Retry-After", "60")
writeJSON(w, http.StatusTooManyRequests, errorResponse{
Error: "too many requests, try again shortly",
RequestID: requestIDFrom(r.Context()),
})
return
}
next.ServeHTTP(w, r)
})
}
@@ -51,14 +51,14 @@ func NewServer(log *slog.Logger, uc UseCases, issuer auth.Issuer) http.Handler {
// Method-aware patterns (Go 1.22+): a GET to /register is a 405, not a match.
mux.HandleFunc("GET /health", h.handleHealth)
mux.HandleFunc("POST /register", h.handleRegister)
mux.HandleFunc("POST /login", h.handleLogin)
mux.Handle("POST /login", rateLimited(newIPLimiter(loginRatePerMinute, loginBurst), http.HandlerFunc(h.handleLogin)))
// /me proves a token round-trips; it sits behind JWT auth.
mux.Handle("GET /me", chain(http.HandlerFunc(h.handleMe), withJWT(issuer)))
// Worker keys: a user mints a long-lived key (JWT-protected), and a worker
// trades it for a short-lived JWT on the public exchange endpoint — the key
// itself is the credential there, so no prior token is required.
mux.HandleFunc("POST /worker-tokens/exchange", h.handleExchangeWorkerKey)
mux.Handle("POST /worker-tokens/exchange", rateLimited(newIPLimiter(exchangeRatePerMinute, exchangeBurst), http.HandlerFunc(h.handleExchangeWorkerKey)))
mux.Handle("POST /worker-keys", chain(http.HandlerFunc(h.handleCreateWorkerKey), withJWT(issuer)))
mux.Handle("GET /worker-keys", chain(http.HandlerFunc(h.handleListWorkerKeys), withJWT(issuer)))
mux.Handle("DELETE /worker-keys/{id}", chain(http.HandlerFunc(h.handleRevokeWorkerKey), withJWT(issuer)))
@@ -597,6 +597,206 @@
"upload_ready": true,
"verifier": "exact-artifact@1",
"version": "1.0.0"
},
{
"capabilities": [
"similarity-search-parallel"
],
"description": "Exact top-k Tanimoto molecular similarity search over deterministic TSV shards with a bounded merge; each shard is fingerprinted and scored across a thread pool. Output is byte-identical to similarity-search.",
"determinism": "byte_exact",
"enabled": true,
"inputs": {
"input": {
"allow_nested_collections": false,
"canonicalizer": "scimesh-tsv-v1",
"encoding": "utf-8",
"max_bytes": 10737418240,
"max_dimensions": [],
"max_records": 100000000,
"media_type": "text/tab-separated-values",
"privacy_class": "project",
"ref": "molecule-table@1",
"retention_class": "durable",
"streaming": false,
"validator": "delimited-table@1",
"validator_configuration": {
"required_columns": [
"canonical_smiles",
"chembl_id"
]
}
}
},
"name": "similarity-search-parallel",
"outputs": {
"result": {
"allow_nested_collections": false,
"canonicalizer": "scimesh-search-result-v1",
"encoding": "utf-8",
"max_bytes": 1073741824,
"max_dimensions": [],
"max_records": 100000,
"media_type": "text/csv",
"privacy_class": "project",
"ref": "similarity-search-result@1",
"retention_class": "durable",
"streaming": false,
"validator": "delimited-table@1",
"validator_configuration": {
"columns": [
"rank",
"chembl_id",
"canonical_smiles",
"similarity"
]
}
}
},
"parameters_schema": {
"additionalProperties": false,
"oneOf": [
{
"not": {
"required": [
"query_smiles"
]
},
"required": [
"query_id"
]
},
{
"not": {
"required": [
"query_id"
]
},
"required": [
"query_smiles"
]
}
],
"properties": {
"max_rows": {
"minimum": 1,
"type": "integer"
},
"progress_every": {
"minimum": 0,
"type": "integer"
},
"query_id": {
"maxLength": 200,
"minLength": 1,
"type": "string"
},
"query_smiles": {
"maxLength": 200,
"minLength": 1,
"type": "string"
},
"threads": {
"description": "Threads used to fingerprint and score one shard (default: CPU count).",
"minimum": 1,
"type": "integer"
},
"threshold": {
"maximum": 1,
"minimum": 0,
"type": "number"
},
"threshold_direction": {
"enum": [
"greater",
"less"
]
},
"top_k": {
"minimum": 1,
"type": "integer"
}
},
"type": "object"
},
"reduction": "top-k",
"trust_modes": [
"trusted",
"untrusted_quorum"
],
"ui_elements": [
{
"default": null,
"field": "query_id",
"group": "",
"help": "ChEMBL id of the query molecule. Provide exactly one of id or SMILES.",
"label": "Query molecule id",
"options": [],
"order": 1,
"placeholder": "",
"widget": "text"
},
{
"default": null,
"field": "query_smiles",
"group": "",
"help": "SMILES of the query molecule. Provide exactly one of id or SMILES.",
"label": "Query molecule SMILES",
"options": [],
"order": 2,
"placeholder": "",
"widget": "text"
},
{
"default": 20,
"field": "top_k",
"group": "",
"help": "Number of most similar molecules to keep per shard (global merge keeps the best of these).",
"label": "Top k",
"options": [],
"order": 3,
"placeholder": "",
"widget": "number"
},
{
"default": "greater",
"field": "threshold_direction",
"group": "",
"help": "Keep molecules with similarity greater or less than the threshold.",
"label": "Direction",
"options": [
"greater",
"less"
],
"order": 4,
"placeholder": "",
"widget": "select"
},
{
"default": null,
"field": "threshold",
"group": "",
"help": "Optional similarity bound: results are filtered to this direction.",
"label": "Similarity threshold",
"options": [],
"order": 5,
"placeholder": "e.g. 0.8",
"widget": "number"
},
{
"default": null,
"field": "threads",
"group": "",
"help": "Threads used to fingerprint and score one shard (default: CPU count).",
"label": "Threads per shard",
"options": [],
"order": 6,
"placeholder": "auto",
"widget": "number"
}
],
"upload_ready": true,
"verifier": "exact-artifact@1",
"version": "1.0.0"
}
]
}
+25
View File
@@ -57,6 +57,31 @@ $Target = Join-Path $InstallDir "$Binary.exe"
Write-Host "Downloading $Url"
Invoke-WebRequest -Uri $Url -OutFile "$Target.tmp"
# Verify the SHA-256 checksum from the release before installing (see
# install.sh for the caveats). $env:SCIMESH_SKIP_VERIFY -eq "1" bypasses.
if ($env:SCIMESH_SKIP_VERIFY -ne "1") {
try {
$SumUrl = "https://github.com/$Repo/releases/download/$Version/SHA256SUMS.txt"
$Sums = (Invoke-WebRequest -Uri $SumUrl).Content
$BinaryName = Split-Path $Url -Leaf
$Line = ($Sums -split "`n") | Where-Object { $_.Trim().EndsWith(" " + $BinaryName) } | Select-Object -First 1
if ($Line) {
$Expected = ($Line -split "\s+")[0]
$Actual = (Get-FileHash -Algorithm SHA256 -Path "$Target.tmp").Hash.ToLower()
if ($Actual -ne $Expected.ToLower()) {
Remove-Item -Force "$Target.tmp"
throw "checksum mismatch for $Binary (got $Actual, want $Expected)"
}
Write-Host "Checksum verified ($($Expected.Substring(0,12))...)"
} else {
Write-Host "WARNING: no checksum entry for $Binary; skipping verification"
}
} catch {
Write-Host "WARNING: could not verify checksum ($($_.Exception.Message)); continuing"
}
}
Move-Item -Force "$Target.tmp" $Target
Write-Host ""
+27
View File
@@ -60,6 +60,33 @@ TARGET="$INSTALL_DIR/$BINARY"
URL="https://github.com/${REPO}/releases/download/${VERSION}/${BINARY}-${OS}-${ARCH}"
echo "Downloading $URL"
curl -fsSL -o "$TARGET.tmp" "$URL"
# Verify the SHA-256 checksum from the release before installing. This guards
# against corrupted downloads and stale CDN caches; it does not protect
# against an active MITM on the same channel (the checksum file travels it
# too). Set SCIMESH_SKIP_VERIFY=1 to bypass.
if [ "${SCIMESH_SKIP_VERIFY:-0}" != "1" ]; then
if SUMFILE=$(mktemp) && curl -fsSL -o "$SUMFILE" "https://github.com/${REPO}/releases/download/${VERSION}/SHA256SUMS.txt"; then
EXPECTED=$(awk '$2 == "'"$(basename "$URL")"'" {print $1}' "$SUMFILE" 2>/dev/null | head -1)
rm -f "$SUMFILE"
if [ -n "$EXPECTED" ]; then
ACTUAL=$(sha256sum "$TARGET.tmp" | awk '{print $1}')
if [ "$ACTUAL" != "$EXPECTED" ]; then
rm -f "$TARGET.tmp"
echo "ERROR: checksum mismatch for $BINARY (got $ACTUAL, want $EXPECTED)" >&2
echo "The download may be corrupted or served by a stale cache. Retry later, or" >&2
echo "pin the version with SCIMESH_VERSION=${VERSION} and re-run." >&2
exit 1
fi
echo "Checksum verified ($(echo "$EXPECTED" | cut -c1-12)…)"
else
echo "WARNING: no checksum entry for $(basename "$URL"); skipping verification"
fi
else
echo "WARNING: could not fetch SHA256SUMS.txt; skipping verification"
fi
fi
chmod +x "$TARGET.tmp"
mv "$TARGET.tmp" "$TARGET"
+1
View File
@@ -23,6 +23,7 @@ scimesh = "scimesh.cli:main"
[project.entry-points."scimesh.workloads"]
"similarity-search@1.0.0" = "scimesh.workloads.search:workload_definition"
"similarity-search-parallel@1.0.0" = "scimesh.workloads.search_parallel:workload_definition"
"similarity-graph@1.0.0" = "scimesh.workloads.graph:workload_definition"
"descriptor-batch@1.0.0" = "scimesh.workloads.descriptors:workload_definition"
"molwt-filter@1.0.0" = "scimesh.workloads.molwt_filter:workload_definition"
+6
View File
@@ -21,6 +21,7 @@ from .environment import current_environment_digest
from .graph import similarity_graph_sdk_definition
from .molwt_filter import molwt_filter_sdk_definition
from .search import similarity_search_sdk_definition
from .search_parallel import similarity_search_parallel_sdk_definition
__all__ = [
"default_sdk_registry",
@@ -46,6 +47,10 @@ def default_sdk_registry(
similarity_search_sdk_definition(shard_rows=shard_rows).definition(),
enabled=True,
)
registry.register(
similarity_search_parallel_sdk_definition(shard_rows=shard_rows).definition(),
enabled=True,
)
registry.register(
similarity_graph_sdk_definition().definition(),
enabled=True,
@@ -82,6 +87,7 @@ def default_sdk_runtime(
workload_capabilities
or (
"similarity-search",
"similarity-search-parallel",
"similarity-graph",
"descriptor-batch",
"molwt-filter",
@@ -0,0 +1,28 @@
"""SDK-built ``similarity-search-parallel`` workload.
Same contract as ``similarity-search`` with a per-shard thread pool. See
``core.py`` for the parallel scoring core and ``definition.py`` for the
manifest-backed handlers.
"""
from .core import (
run_search_shard_parallel,
search_similar_parallel,
write_search_shards,
)
from .definition import (
MAP_ENTRY_POINT,
SimilaritySearchParallelSDKWorkload,
similarity_search_parallel_sdk_definition,
workload_definition,
)
__all__ = [
"MAP_ENTRY_POINT",
"SimilaritySearchParallelSDKWorkload",
"similarity_search_parallel_sdk_definition",
"workload_definition",
"run_search_shard_parallel",
"search_similar_parallel",
"write_search_shards",
]
+199
View File
@@ -0,0 +1,199 @@
"""Scientific core for the SDK-built ``similarity-search-parallel`` workload.
The exact same semantics as ``similarity-search`` identical partial format,
identical bounded merge, byte-identical output but the per-molecule
fingerprinting and Tanimoto scoring of one shard run across a thread pool
(``threads`` parameter, default = CPU count).
Parallelism is confined to the scoring phase: ``ThreadPoolExecutor.map`` keeps
the input row order, so the results are merged exactly like the sequential
reference (same ``_HeapEntry`` logic), which makes the output byte-identical
for every thread count by construction.
"""
from __future__ import annotations
import heapq
import os
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Mapping
from rdkit import Chem
from rdkit.Chem import DataStructs
from scimesh.chemistry.dataset import MoleculeRecord, parse_smiles
from scimesh.chemistry.fingerprints import fingerprint
from scimesh.workloads.search.core import (
run_search_shard,
write_search_partial,
write_search_shards,
)
from scimesh.workloads.similarity_search import (
DatasetStats,
SearchResult,
SimilarityMatch,
_HeapEntry,
iter_valid_molecules,
)
def search_similar_parallel(
tsv_path: Path,
query: MoleculeRecord,
top_k: int,
*,
threads: int = 0,
max_rows: int | None = None,
threshold: float | None = None,
threshold_direction: str = "greater",
) -> SearchResult:
"""Exact top-k matches with a bounded heap, scored by a thread pool.
Identical selection and ordering to ``search_similar`` for every thread
count: the merge runs in row order over the parallel-computed scores.
"""
if top_k < 1:
raise ValueError("--top-k must be a positive integer")
if threads < 0:
raise ValueError("threads must be a non-negative integer")
if threshold is not None and not 0.0 <= threshold <= 1.0:
raise ValueError("--threshold must be between 0 and 1")
if threshold_direction not in {"greater", "less"}:
raise ValueError("--threshold-direction must be 'greater' or 'less'")
workers = threads or (os.cpu_count() or 1)
query_fingerprint = fingerprint(query.molecule)
query_canonical_smiles = Chem.MolToSmiles(query.molecule, canonical=True)
stats = DatasetStats()
records = list(iter_valid_molecules(tsv_path, stats, max_rows=max_rows))
def score(record: MoleculeRecord):
candidate_smiles = Chem.MolToSmiles(record.molecule, canonical=True)
if (
record.molecule_id == query.molecule_id
or candidate_smiles == query_canonical_smiles
):
return None
similarity = DataStructs.TanimotoSimilarity(
query_fingerprint, fingerprint(record.molecule)
)
if threshold is not None and (
similarity < threshold
if threshold_direction == "greater"
else similarity > threshold
):
return None
return similarity
# map preserves the input order, so the merge below is exactly the
# sequential reference's merge, just over precomputed scores.
with ThreadPoolExecutor(max_workers=workers) as pool:
scored = pool.map(score, records)
heap: list[_HeapEntry] = []
for record, similarity in zip(records, scored):
if similarity is None:
continue
match = SimilarityMatch(similarity, record.molecule_id, record.smiles)
rank_key = match.sort_key(threshold_direction)
entry = _HeapEntry(match, rank_key)
if len(heap) < top_k:
heapq.heappush(heap, entry)
elif rank_key < heap[0].rank_key:
heapq.heapreplace(heap, entry)
matches = [entry.match for entry in sorted(heap, key=lambda e: e.rank_key)]
return SearchResult(matches=matches, stats=stats)
def run_search_shard_parallel(
input_path: Path,
parameters: Mapping[str, object],
output_path: Path,
) -> dict[str, int]:
"""Run one planned shard with the parallel scoring core.
Accepts the same parameters as ``similarity-search`` plus ``threads``.
"""
allowed = {
"query_id",
"query_smiles",
"top_k",
"threshold",
"threshold_direction",
"progress_every",
"threads",
}
unknown = set(parameters) - allowed
if unknown:
raise ValueError(
f"unsupported similarity-search-parallel parameters: {', '.join(sorted(unknown))}"
)
query_smiles = parameters.get("query_smiles")
query_id = parameters.get("query_id")
if isinstance(query_id, str) and not isinstance(query_smiles, str):
from rdkit import Chem
from scimesh.chemistry.dataset import find_molecule_by_id
record = find_molecule_by_id(input_path, query_id)
query_smiles = Chem.MolToSmiles(record.molecule, canonical=True)
if not isinstance(query_smiles, str) or not query_smiles.strip():
raise ValueError("query_smiles is required for a distributed shard")
molecule = parse_smiles(query_smiles)
if molecule is None:
raise ValueError("query_smiles is invalid")
top_k = _positive_int(parameters.get("top_k", 20), "top_k")
threads = _nonnegative_int(parameters.get("threads", 0), "threads")
threshold = None
if "threshold" in parameters:
threshold = _unit_interval(parameters["threshold"], "threshold")
direction = parameters.get("threshold_direction", "greater")
if direction not in {"greater", "less"}:
raise ValueError("threshold_direction must be 'greater' or 'less'")
assert isinstance(direction, str)
result = search_similar_parallel(
input_path,
MoleculeRecord("query", query_smiles, molecule),
top_k=top_k,
threads=threads,
threshold=threshold,
threshold_direction=direction,
)
write_search_partial(output_path, result.matches)
return {
"scanned_rows": result.stats.scanned,
"valid_molecules": result.stats.valid,
"invalid_smiles": result.stats.invalid,
"matches_emitted": len(result.matches),
}
def _positive_int(value: object, name: str) -> int:
if isinstance(value, bool) or not isinstance(value, int) or value < 1:
raise ValueError(f"{name} must be a positive integer")
return value
def _nonnegative_int(value: object, name: str) -> int:
if isinstance(value, bool) or not isinstance(value, int) or value < 0:
raise ValueError(f"{name} must be a non-negative integer")
return value
def _unit_interval(value: object, name: str) -> float:
if isinstance(value, bool) or not isinstance(value, (int, float)):
raise ValueError(f"{name} must be a number between 0 and 1")
return float(value)
# The shared shard writer is re-exported so the workload definition can reuse
# the deterministic partitioning without importing search internals.
__all__ = [
"search_similar_parallel",
"run_search_shard_parallel",
"write_search_shards",
"run_search_shard",
]
@@ -0,0 +1,144 @@
"""SDK-built ``similarity-search-parallel`` workload definition and handlers.
A subclass of ``SimilaritySearchSDKWorkload``: identical contract (plan-time
query resolution, deterministic sharding, top-k reduction, byte-identical
partials), but each shard's fingerprinting and scoring runs across a thread
pool (``threads`` parameter, default = CPU count).
"""
from __future__ import annotations
from collections.abc import Mapping
from pathlib import Path
from typing import Any
from scimesh.sdk.batch import MapReduceWorkload
from scimesh.sdk.identity import WorkloadId
from scimesh.sdk.plans import JobRequest
from scimesh.sdk.registry import WorkloadDefinition
from scimesh.sdk.ui import UIElement
from ..environment import current_environment_digest, current_scimesh_package_digest
from ..search.definition import SimilaritySearchSDKWorkload, _parameters_schema
from .core import run_search_shard_parallel
MAP_ENTRY_POINT = "scimesh.workloads.search_parallel.definition:map_search_parallel@v1"
# The parallel variant adds only the thread-count parameter on top of the
# search contract; everything else (schemas, entry points of the reduce stage,
# partitioning) is inherited.
_MAP_PARAMETERS = (
"query_id",
"query_smiles",
"top_k",
"threshold",
"threshold_direction",
"progress_every",
"threads",
)
def _parallel_parameters_schema() -> dict[str, Any]:
schema = dict(_parameters_schema())
properties = dict(schema["properties"])
properties["threads"] = {
"type": "integer",
"minimum": 1,
"description": "Threads used to fingerprint and score one shard (default: CPU count).",
}
schema["properties"] = properties
return schema
class SimilaritySearchParallelSDKWorkload(SimilaritySearchSDKWorkload):
"""Exact top-k Tanimoto search with a per-shard thread pool."""
workload_id = WorkloadId("similarity-search-parallel", "1.0.0")
description = (
"Exact top-k Tanimoto molecular similarity search over deterministic "
"TSV shards with a bounded merge; each shard is fingerprinted and "
"scored across a thread pool. Output is byte-identical to "
"similarity-search."
)
parameters_schema = _parallel_parameters_schema()
map_parameter_names = _MAP_PARAMETERS
map_entry_point = MAP_ENTRY_POINT
ui_elements = SimilaritySearchSDKWorkload.ui_elements + (
UIElement(
"threads",
"number",
"Threads per shard",
help="Threads used to fingerprint and score one shard (default: CPU count).",
placeholder="auto",
order=6,
),
)
def domain_validate(self, parameters: Mapping[str, Any]) -> None:
# The search base rejects unknown parameters; threads is our addition,
# so it is validated here and stripped before delegating.
rest = dict(parameters)
threads = rest.pop("threads", None)
if threads is not None and (
isinstance(threads, bool) or not isinstance(threads, int) or threads < 1
):
raise ValueError("threads must be a positive integer")
super().domain_validate(rest)
def resolved_parameters_for_plan(
self,
job,
input_path,
resolved,
):
# threads is a map-stage-only knob; strip it from the plan-level
# resolved parameters so the reduce stage projection stays clean.
resolved = super().resolved_parameters_for_plan(job, input_path, resolved)
stripped = dict(resolved)
stripped.pop("threads", None)
return stripped
def resolved_parameters(self, request: JobRequest) -> dict[str, Any]:
resolved = super().resolved_parameters(request)
if "threads" in request.parameters:
threads = request.parameters["threads"]
if isinstance(threads, bool) or not isinstance(threads, int) or threads < 1:
raise ValueError("threads must be a positive integer")
resolved["threads"] = threads
return resolved
def compute_shard(
self,
inputs: Mapping[str, Path],
parameters: Mapping[str, Any],
output_path: Path,
) -> Mapping[str, int | float]:
return run_search_shard_parallel(inputs["input"], parameters, output_path)
def similarity_search_parallel_sdk_definition(
*,
shard_rows: int = 10_000,
package_digest: str | None = None,
environment_digest: str | None = None,
) -> SimilaritySearchParallelSDKWorkload:
"""Build the SDK-built parallel similarity-search definition for tests."""
return SimilaritySearchParallelSDKWorkload(
shard_rows=shard_rows,
package_digest=package_digest or current_scimesh_package_digest(),
environment_digest=environment_digest or current_environment_digest(),
)
def workload_definition() -> WorkloadDefinition:
"""Installed entry-point factory for the SDK-built parallel search."""
return similarity_search_parallel_sdk_definition().definition()
def map_search_parallel(
input_path: Path,
parameters: Mapping[str, object],
output_path: Path,
) -> dict[str, int]:
"""Digest-pinned map entry point for the parallel search shard."""
return run_search_shard_parallel(input_path, parameters, output_path)
+7 -1
View File
@@ -244,7 +244,13 @@ def test_workload_cli_exports_the_library_as_json(tmp_path: Path) -> None:
assert payload["schema_version"] == 2
names = [item["name"] for item in payload["workloads"]]
assert names == sorted(
["descriptor-batch", "molwt-filter", "similarity-graph", "similarity-search"]
[
"descriptor-batch",
"molwt-filter",
"similarity-graph",
"similarity-search",
"similarity-search-parallel",
]
)
for item in payload["workloads"]:
assert item["version"] == "1.0.0"
+1 -1
View File
@@ -60,7 +60,7 @@ def _registered_similarity_search(shard_rows: int = 2):
registry = default_sdk_registry(shard_rows=shard_rows)
runtime = default_sdk_runtime()
descriptions = registry.descriptions()
assert len(descriptions) == 4
assert len(descriptions) == 5
description = next(
item for item in descriptions if item.workload.name == "similarity-search"
)
+185
View File
@@ -0,0 +1,185 @@
"""Tests for the SDK-built similarity-search-parallel workload."""
from __future__ import annotations
import csv
from pathlib import Path
import pytest
from scimesh.sdk import (
ArtifactCollection,
DeterminismProfile,
JobRequest,
LocalArtifactStore,
LocalCoreBatchExecutor,
LocalPlanningContext,
StageKind,
)
from scimesh.workloads.library import default_sdk_registry, default_sdk_runtime
from scimesh.workloads.search.core import run_search_shard, write_search_shards
from scimesh.workloads.search_parallel import (
run_search_shard_parallel,
search_similar_parallel,
)
from scimesh.workloads.similarity_search import (
find_molecule_by_id,
search_similar,
write_search_results,
)
def _write_dataset(path: Path, molecules: list[tuple[str, str]]) -> None:
path.write_text(
"chembl_id\tcanonical_smiles\n"
+ "".join(f"{mid}\t{smiles}\n" for mid, smiles in molecules),
encoding="utf-8",
)
def _tie_dataset(path: Path) -> None:
# Deliberate similarity ties: propanol isomers and duplicated rows, so the
# parallel merge must reproduce the sequential row-order preference.
_write_dataset(
path,
[
("QUERY", "CCO"),
("A1", "CCCO"),
("A2", "C(CC)O"),
("B", "CCN"),
("C1", "CCC"),
("C2", "CCC"),
("D", "CCCC"),
],
)
def test_parallel_matches_sequential_byte_exactly(tmp_path: Path) -> None:
dataset = tmp_path / "molecules.tsv"
_tie_dataset(dataset)
query = find_molecule_by_id(dataset, "QUERY")
reference = search_similar(dataset, query, top_k=5, progress_every=0)
reference_path = tmp_path / "reference.csv"
write_search_results(reference_path, reference.matches)
for threads in (1, 2, 4):
parallel = search_similar_parallel(dataset, query, top_k=5, threads=threads)
parallel_path = tmp_path / f"parallel-{threads}.csv"
write_search_results(parallel_path, parallel.matches)
assert parallel_path.read_bytes() == reference_path.read_bytes(), (
f"threads={threads} diverged from the reference"
)
def test_parallel_shard_matches_sequential_shard(tmp_path: Path) -> None:
dataset = tmp_path / "molecules.tsv"
_tie_dataset(dataset)
shard_dir = tmp_path / "shards"
shard_dir.mkdir()
shards = write_search_shards(dataset, shard_dir, shard_rows=2)
parameters = {"query_smiles": "CCO", "top_k": 3, "threads": 4}
sequential_out = tmp_path / "seq.tsv"
parallel_out = tmp_path / "par.tsv"
sequential_parameters = dict(parameters)
sequential_parameters.pop("threads")
run_search_shard(shards[0], sequential_parameters, sequential_out)
run_search_shard_parallel(shards[0], parameters, parallel_out)
assert parallel_out.read_bytes() == sequential_out.read_bytes()
with parallel_out.open(encoding="utf-8") as handle:
rows = list(csv.DictReader(handle))
assert rows[0]["rank"] == "1"
assert rows[0]["similarity"].startswith("0.5") # CCO vs CCCO
assert len(rows) <= 3
def test_parallel_rejects_bad_parameters(tmp_path: Path) -> None:
dataset = tmp_path / "molecules.tsv"
_tie_dataset(dataset)
output = tmp_path / "out.tsv"
with pytest.raises(ValueError, match="threads must be a non-negative integer"):
run_search_shard_parallel(
dataset, {"query_smiles": "CCO", "threads": -1}, output
)
with pytest.raises(ValueError, match="unsupported"):
run_search_shard_parallel(dataset, {"query_smiles": "CCO", "nope": 1}, output)
with pytest.raises(ValueError, match="query_smiles is invalid"):
run_search_shard_parallel(dataset, {"query_smiles": "СС"}, output)
def _registered_parallel_search(shard_rows: int = 2):
registry = default_sdk_registry(shard_rows=shard_rows)
runtime = default_sdk_runtime()
description = next(
item
for item in registry.descriptions()
if item.workload.name == "similarity-search-parallel"
)
definition, negotiated = registry.require(
description.workload.name,
description.workload.version,
description.package_digest,
runtime=runtime,
)
return registry, runtime, description, definition, negotiated
def test_parallel_manifest_is_registered_and_negotiable() -> None:
_, runtime, description, definition, negotiated = _registered_parallel_search()
manifest = definition.manifest
assert description.enabled is True
assert manifest.workload.name == "similarity-search-parallel"
assert manifest.workload.version == "1.0.0"
assert manifest.determinism is DeterminismProfile.BYTE_EXACT
assert manifest.verifier.verifier.canonical == "exact-artifact@1"
assert set(mode.value for mode in manifest.trust_modes) == {
"trusted",
"untrusted_quorum",
}
assert [stage.kind for stage in manifest.workflow.stages] == [
StageKind.MAP,
StageKind.REDUCE,
]
assert "threads" in manifest.parameters_schema["properties"]
assert negotiated is not None
assert runtime is not None
def test_parallel_executor_matches_reference(tmp_path: Path) -> None:
dataset = tmp_path / "molecules.tsv"
_tie_dataset(dataset)
registry, runtime, description, definition, _ = _registered_parallel_search()
artifact_store = LocalArtifactStore(tmp_path / "artifacts")
input_port = definition.manifest.inputs["input"]
dataset_artifact = artifact_store.import_file(
dataset,
declaration=input_port.schema,
)
request = JobRequest(
workload=definition.manifest.workload,
parameters={"query_id": "QUERY", "top_k": 3, "threads": 2},
inputs={"input": ArtifactCollection.single(dataset_artifact)},
)
result = LocalCoreBatchExecutor(
registry,
runtime,
artifact_store,
tmp_path / "sdk-work",
).execute(request, description.package_digest)
result_artifact = result.outputs["result"].items[0].artifact
reference_path = tmp_path / "reference.csv"
query = find_molecule_by_id(dataset, "QUERY")
reference = search_similar(dataset, query, top_k=3, progress_every=0)
write_search_results(reference_path, reference.matches)
assert (
artifact_store.materialize(result_artifact).read_bytes()
== reference_path.read_bytes()
)
assert result.task_key == "reduce/final"
+66
View File
@@ -0,0 +1,66 @@
# SciMesh uninstaller for Windows: stops the running component, removes its
# binary, and — when asked — deletes its data directory.
#
# powershell -ExecutionPolicy Bypass -c "irm https://raw.githubusercontent.com/emil28092005/SciMesh/main/uninstall.ps1 | iex"
#
# $env:SCIMESH_COMPONENT selects the component (coordinator | worker | all,
# default all). Data is kept unless $env:SCIMESH_PURGE -eq "1" or the
# -Purge switch is passed (in a pipe there is no interactive prompt).
param(
[switch]$Purge
)
$ErrorActionPreference = "Stop"
$Component = if ($env:SCIMESH_COMPONENT) { $env:SCIMESH_COMPONENT } else { "all" }
if ($Purge -or $env:SCIMESH_PURGE -eq "1") { $Purge = $true } else { $Purge = $false }
$InstallDir = if ($env:SCIMESH_INSTALL_DIR) {
$env:SCIMESH_INSTALL_DIR
} else {
Join-Path $env:LOCALAPPDATA "SciMesh"
}
function Remove-Component {
param([string]$Name, [string]$Binary)
Write-Host "Stopping $Name"
Get-Process -Name $Name -ErrorAction SilentlyContinue | Stop-Process -Force -ErrorAction SilentlyContinue
$target = Join-Path $InstallDir $Binary
if (Test-Path $target) {
Write-Host "Removing $target"
Remove-Item -Force $target
}
}
function Remove-Data {
param([string]$Dir, [string]$Label)
if (-not (Test-Path $Dir)) { return }
if ($Purge) {
Remove-Item -Recurse -Force $Dir
Write-Host "Deleted $Dir"
return
}
# No interactive prompt in a pipe: keep the data by default.
Write-Host "Keeping $Dir ($Label). Pass -Purge to delete it."
}
switch ($Component) {
"coordinator" {
Remove-Component -Name "coordinator" -Binary "coordinator.exe"
Remove-Data -Dir (Join-Path $HOME ".scimesh") -Label "secrets, databases, artifacts, users"
}
"worker" {
Remove-Component -Name "worker-agent" -Binary "worker-agent.exe"
Remove-Data -Dir (Join-Path $HOME ".scimesh-worker") -Label "worker config, runtime venv, logs"
}
"all" {
Remove-Component -Name "coordinator" -Binary "coordinator.exe"
Remove-Component -Name "worker-agent" -Binary "worker-agent.exe"
Remove-Data -Dir (Join-Path $HOME ".scimesh") -Label "secrets, databases, artifacts, users"
Remove-Data -Dir (Join-Path $HOME ".scimesh-worker") -Label "worker config, runtime venv, logs"
}
default { throw "unknown component: $Component (use 'coordinator', 'worker' or 'all')" }
}
Write-Host ""
Write-Host "SciMesh $Component uninstalled."
Write-Host "Pass -Purge to also delete the data directories."
+83
View File
@@ -0,0 +1,83 @@
#!/usr/bin/env bash
# SciMesh uninstaller: stops the running component, removes its binary, and —
# when asked — deletes its data directory.
#
# curl -fsSL https://raw.githubusercontent.com/emil28092005/SciMesh/main/uninstall.sh | bash -s coordinator
# curl -fsSL https://raw.githubusercontent.com/emil28092005/SciMesh/main/uninstall.sh | bash -s worker
# curl -fsSL https://raw.githubusercontent.com/emil28092005/SciMesh/main/uninstall.sh # both
#
# Data (jobs, artifacts, users, secrets, the managed venv) is kept by default.
# Pass --purge (or set SCIMESH_PURGE=1) to delete it without asking; without
# it the script prompts interactively. Because `curl | bash` pipes have no
# interactive stdin, the default is always "keep data".
set -eu
COMPONENT="${1:-all}"
PURGE=0
for arg in "$@"; do
case "$arg" in
--purge) PURGE=1 ;;
esac
done
[ "${SCIMESH_PURGE:-0}" = "1" ] && PURGE=1
INSTALL_DIR="${SCIMESH_INSTALL_DIR:-$HOME/.local/bin}"
remove_component() {
case "$1" in
coordinator)
echo "Stopping the coordinator (serve)…"
pkill -x coordinator 2>/dev/null || true
echo "Removing $INSTALL_DIR/coordinator…"
rm -f "$INSTALL_DIR/coordinator"
;;
worker)
echo "Stopping the worker agent and its setup wizard…"
pkill -x worker-agent 2>/dev/null || true
echo "Removing $INSTALL_DIR/worker-agent…"
rm -f "$INSTALL_DIR/worker-agent"
;;
*) echo "unknown component: $1 (use 'coordinator', 'worker' or 'all')" >&2; exit 1 ;;
esac
}
remove_data() {
local dir="$1" label="$2"
if [ ! -d "$dir" ]; then
return 0
fi
if [ "$PURGE" = "1" ]; then
rm -rf "$dir"
echo "Deleted $dir"
return 0
fi
# In a pipe (curl | bash) stdin is exhausted, so the prompt defaults to keep.
printf "Delete %s (%s)? [y/N] " "$dir" "$label"
read -r answer || answer=""
case "$answer" in
y|Y|yes|YES) rm -rf "$dir"; echo "Deleted $dir" ;;
*) echo "Keeping $dir" ;;
esac
}
case "$COMPONENT" in
coordinator)
remove_component coordinator
remove_data "$HOME/.scimesh" "secrets, databases, artifacts, users, managed venv"
;;
worker)
remove_component worker
remove_data "$HOME/.scimesh-worker" "worker config, runtime venv, logs"
;;
all)
remove_component coordinator
remove_component worker
remove_data "$HOME/.scimesh" "secrets, databases, artifacts, users, managed venv"
remove_data "$HOME/.scimesh-worker" "worker config, runtime venv, logs"
;;
*) echo "unknown component: $COMPONENT (use 'coordinator', 'worker' or 'all')" >&2; exit 1 ;;
esac
echo
echo "SciMesh $COMPONENT uninstalled."
echo "Pass --purge to also delete the data directories without asking."