From 4c792400c1c66cde166a7393286a1b629e853e65 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Boschi?= Date: Tue, 20 Jan 2026 10:17:56 +0100 Subject: [PATCH] feat: new 'worker' service (#176) * feat: new 'worker' service * doc * docs * tests --- helm/hindsight/templates/_helpers.tpl | 16 + helm/hindsight/templates/api-deployment.yaml | 5 + helm/hindsight/templates/worker-service.yaml | 25 + .../templates/worker-statefulset.yaml | 110 ++++ helm/hindsight/values.yaml | 57 ++ hindsight-api/hindsight_api/admin/cli.py | 59 ++ .../l7g8h9i0j1k2_add_worker_columns.py | 109 ++++ hindsight-api/hindsight_api/api/http.py | 35 ++ hindsight-api/hindsight_api/config.py | 48 +- .../hindsight_api/engine/memory_engine.py | 18 +- .../hindsight_api/engine/task_backend.py | 302 ++++----- hindsight-api/hindsight_api/main.py | 9 +- .../hindsight_api/worker/__init__.py | 11 + hindsight-api/hindsight_api/worker/main.py | 285 +++++++++ hindsight-api/hindsight_api/worker/poller.py | 281 +++++++++ hindsight-api/pyproject.toml | 1 + hindsight-api/tests/conftest.py | 3 + .../tests/test_custom_embedding_dimension.py | 4 + hindsight-api/tests/test_load_large_batch.py | 2 + hindsight-api/tests/test_worker.py | 592 ++++++++++++++++++ hindsight-docs/docs/developer/admin-cli.md | 47 ++ .../docs/developer/configuration.md | 13 +- hindsight-docs/docs/developer/installation.md | 16 +- hindsight-docs/docs/developer/services.md | 50 +- scripts/dev/start-worker.sh | 20 + 25 files changed, 1866 insertions(+), 252 deletions(-) create mode 100644 helm/hindsight/templates/worker-service.yaml create mode 100644 helm/hindsight/templates/worker-statefulset.yaml create mode 100644 hindsight-api/hindsight_api/alembic/versions/l7g8h9i0j1k2_add_worker_columns.py create mode 100644 hindsight-api/hindsight_api/worker/__init__.py create mode 100644 hindsight-api/hindsight_api/worker/main.py create mode 100644 hindsight-api/hindsight_api/worker/poller.py create mode 100644 hindsight-api/tests/test_worker.py create mode 100755 scripts/dev/start-worker.sh diff --git a/helm/hindsight/templates/_helpers.tpl b/helm/hindsight/templates/_helpers.tpl index 418a41c7..70d8a421 100644 --- a/helm/hindsight/templates/_helpers.tpl +++ b/helm/hindsight/templates/_helpers.tpl @@ -80,6 +80,22 @@ Control plane selector labels app.kubernetes.io/component: control-plane {{- end }} +{{/* +Worker labels +*/}} +{{- define "hindsight.worker.labels" -}} +{{ include "hindsight.labels" . }} +app.kubernetes.io/component: worker +{{- end }} + +{{/* +Worker selector labels +*/}} +{{- define "hindsight.worker.selectorLabels" -}} +{{ include "hindsight.selectorLabels" . }} +app.kubernetes.io/component: worker +{{- end }} + {{/* Create the name of the service account to use */}} diff --git a/helm/hindsight/templates/api-deployment.yaml b/helm/hindsight/templates/api-deployment.yaml index 47fb22c9..ccfb30f7 100644 --- a/helm/hindsight/templates/api-deployment.yaml +++ b/helm/hindsight/templates/api-deployment.yaml @@ -55,6 +55,11 @@ spec: {{- end }} - name: HINDSIGHT_API_DATABASE_URL value: {{ include "hindsight.databaseUrl" . | quote }} + {{- /* Disable internal worker when dedicated workers are enabled */}} + {{- if .Values.worker.enabled }} + - name: HINDSIGHT_API_WORKER_ENABLED + value: "false" + {{- end }} {{- range $key, $value := .Values.api.env }} - name: {{ $key }} value: {{ $value | quote }} diff --git a/helm/hindsight/templates/worker-service.yaml b/helm/hindsight/templates/worker-service.yaml new file mode 100644 index 00000000..1244cf0c --- /dev/null +++ b/helm/hindsight/templates/worker-service.yaml @@ -0,0 +1,25 @@ +{{- if .Values.worker.enabled }} +apiVersion: v1 +kind: Service +metadata: + name: {{ include "hindsight.fullname" . }}-worker + labels: + {{- include "hindsight.worker.labels" . | nindent 4 }} + {{- if .Values.podAnnotations }} + annotations: + {{- /* Common Prometheus annotations for metrics scraping */}} + prometheus.io/scrape: "true" + prometheus.io/port: {{ .Values.worker.service.port | quote }} + prometheus.io/path: "/metrics" + {{- end }} +spec: + # Headless service for StatefulSet (enables stable DNS names like worker-0.worker.namespace) + clusterIP: None + ports: + - port: {{ .Values.worker.service.port }} + targetPort: {{ .Values.worker.service.targetPort }} + protocol: TCP + name: http + selector: + {{- include "hindsight.worker.selectorLabels" . | nindent 4 }} +{{- end }} diff --git a/helm/hindsight/templates/worker-statefulset.yaml b/helm/hindsight/templates/worker-statefulset.yaml new file mode 100644 index 00000000..2894e0b4 --- /dev/null +++ b/helm/hindsight/templates/worker-statefulset.yaml @@ -0,0 +1,110 @@ +{{- if .Values.worker.enabled }} +apiVersion: apps/v1 +kind: StatefulSet +metadata: + name: {{ include "hindsight.fullname" . }}-worker + labels: + {{- include "hindsight.worker.labels" . | nindent 4 }} +spec: + serviceName: {{ include "hindsight.fullname" . }}-worker + replicas: {{ .Values.worker.replicaCount }} + selector: + matchLabels: + {{- include "hindsight.worker.selectorLabels" . | nindent 6 }} + template: + metadata: + annotations: + {{- if not .Values.existingSecret }} + checksum/secret: {{ include (print $.Template.BasePath "/secret.yaml") . | sha256sum }} + {{- end }} + {{- with .Values.podAnnotations }} + {{- toYaml . | nindent 8 }} + {{- end }} + labels: + {{- include "hindsight.worker.selectorLabels" . | nindent 8 }} + spec: + {{- if .Values.serviceAccount.create }} + serviceAccountName: {{ include "hindsight.serviceAccountName" . }} + {{- end }} + securityContext: + {{- toYaml .Values.podSecurityContext | nindent 8 }} + containers: + - name: worker + securityContext: + {{- toYaml .Values.securityContext | nindent 10 }} + image: "{{ .Values.worker.image.repository }}:{{ .Values.worker.image.tag | default .Values.version }}" + imagePullPolicy: {{ .Values.worker.image.pullPolicy }} + command: ["hindsight-worker"] + ports: + - name: http + containerPort: {{ .Values.worker.service.targetPort }} + protocol: TCP + {{- if .Values.existingSecret }} + envFrom: + - secretRef: + name: {{ .Values.existingSecret }} + {{- end }} + env: + {{- /* POSTGRES_PASSWORD must be defined before DATABASE_URL for $(VAR) interpolation */}} + {{- if not .Values.postgresql.enabled }} + - name: POSTGRES_PASSWORD + valueFrom: + secretKeyRef: + name: {{ include "hindsight.secretName" . }} + key: postgres-password + {{- end }} + - name: HINDSIGHT_API_DATABASE_URL + value: {{ include "hindsight.databaseUrl" . | quote }} + {{- /* Worker ID uses pod name (StatefulSet provides stable names like worker-0, worker-1) */}} + - name: HINDSIGHT_API_WORKER_ID + valueFrom: + fieldRef: + fieldPath: metadata.name + {{- /* Inherit LLM config from api.env */}} + {{- range $key, $value := .Values.api.env }} + - name: {{ $key }} + value: {{ $value | quote }} + {{- end }} + {{- /* Worker-specific env vars */}} + {{- range $key, $value := .Values.worker.env }} + - name: {{ $key }} + value: {{ $value | quote }} + {{- end }} + {{- /* Only use secrets when not using existingSecret */}} + {{- if not .Values.existingSecret }} + {{- /* Inherit secrets from api.secrets */}} + {{- range $key, $value := .Values.api.secrets }} + - name: {{ $key }} + valueFrom: + secretKeyRef: + name: {{ include "hindsight.secretName" $ }} + key: {{ $key }} + {{- end }} + {{- /* Worker-specific secrets (can override api.secrets) */}} + {{- range $key, $value := .Values.worker.secrets }} + - name: {{ $key }} + valueFrom: + secretKeyRef: + name: {{ include "hindsight.secretName" $ }} + key: {{ $key }} + {{- end }} + {{- end }} + livenessProbe: + {{- toYaml .Values.worker.livenessProbe | nindent 10 }} + readinessProbe: + {{- toYaml .Values.worker.readinessProbe | nindent 10 }} + resources: + {{- toYaml .Values.worker.resources | nindent 10 }} + {{- with .Values.nodeSelector }} + nodeSelector: + {{- toYaml . | nindent 8 }} + {{- end }} + {{- with .Values.affinity }} + affinity: + {{- toYaml . | nindent 8 }} + {{- end }} + {{- with .Values.tolerations }} + tolerations: + {{- toYaml . | nindent 8 }} + {{- end }} +{{- end }} diff --git a/helm/hindsight/values.yaml b/helm/hindsight/values.yaml index 610aafb4..544f7128 100644 --- a/helm/hindsight/values.yaml +++ b/helm/hindsight/values.yaml @@ -67,6 +67,63 @@ api: # HINDSIGHT_API_LLM_API_KEY: "your-api-key" # HINDSIGHT_API_LLM_BASE_URL: "https://api.groq.com/openai/v1" +# Worker settings (distributed task processing) +# When enabled, dedicated worker pods process tasks and the API's internal worker is disabled +worker: + enabled: false + replicaCount: 2 + image: + repository: ghcr.io/vectorize-io/hindsight-api + pullPolicy: IfNotPresent + # tag defaults to .Values.version if not specified + + service: + # Service for metrics scraping (headless for StatefulSet) + port: 8889 + targetPort: 8889 + + # Resource limits and requests + resources: + limits: + cpu: 2000m + memory: 4Gi + requests: + cpu: 500m + memory: 1Gi + + # Liveness and readiness probes + livenessProbe: + httpGet: + path: /health + port: 8889 + initialDelaySeconds: 30 + periodSeconds: 10 + timeoutSeconds: 5 + failureThreshold: 3 + + readinessProbe: + httpGet: + path: /health + port: 8889 + initialDelaySeconds: 10 + periodSeconds: 5 + timeoutSeconds: 3 + failureThreshold: 3 + + # Worker-specific environment variables + env: + # Poll interval in milliseconds (how often to check for new tasks) + HINDSIGHT_API_WORKER_POLL_INTERVAL_MS: "500" + # Number of tasks to claim per poll cycle + HINDSIGHT_API_WORKER_BATCH_SIZE: "10" + # Max retries before marking a task as failed + HINDSIGHT_API_WORKER_MAX_RETRIES: "3" + # HTTP port for metrics/health (matches service.targetPort) + HINDSIGHT_API_WORKER_HTTP_PORT: "8889" + + # Secret environment variables (inherited from api.secrets if not specified) + secrets: {} + # Image settings for control plane controlPlane: enabled: true diff --git a/hindsight-api/hindsight_api/admin/cli.py b/hindsight-api/hindsight_api/admin/cli.py index ff7b05d5..d5009602 100644 --- a/hindsight-api/hindsight_api/admin/cli.py +++ b/hindsight-api/hindsight_api/admin/cli.py @@ -244,6 +244,65 @@ def run_db_migration( typer.echo("Database migrations completed successfully") +async def _decommission_worker(db_url: str, worker_id: str, schema: str = "public") -> int: + """Release all tasks owned by a worker, setting them back to pending status.""" + is_pg0, instance_name, _ = parse_pg0_url(db_url) + if is_pg0: + typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...") + resolved_url = await resolve_database_url(db_url) + + conn = await asyncpg.connect(resolved_url) + try: + table = _fq_table("async_operations", schema) + result = await conn.fetch( + f""" + UPDATE {table} + SET status = 'pending', worker_id = NULL, claimed_at = NULL, updated_at = now() + WHERE worker_id = $1 AND status = 'processing' + RETURNING operation_id + """, + worker_id, + ) + return len(result) + finally: + await conn.close() + + +@app.command(name="decommission-worker") +def decommission_worker( + worker_id: str = typer.Argument(..., help="Worker ID to decommission"), + schema: str = typer.Option("public", "--schema", "-s", help="Database schema"), + yes: bool = typer.Option(False, "--yes", "-y", help="Skip confirmation prompt"), +): + """Release all tasks owned by a worker (sets status back to pending). + + Use this command when a worker has crashed or been removed without graceful shutdown. + All tasks that were being processed by the worker will be released back to the queue + so other workers can pick them up. + """ + config = HindsightConfig.from_env() + + if not config.database_url: + typer.echo("Error: Database URL not configured.", err=True) + typer.echo("Set HINDSIGHT_API_DATABASE_URL environment variable.", err=True) + raise typer.Exit(1) + + if not yes: + typer.confirm( + f"This will release all tasks owned by worker '{worker_id}' back to pending. Continue?", + abort=True, + ) + + typer.echo(f"Decommissioning worker '{worker_id}' (schema: {schema})...") + + count = asyncio.run(_decommission_worker(config.database_url, worker_id, schema)) + + if count > 0: + typer.echo(f"Released {count} task(s) from worker '{worker_id}'") + else: + typer.echo(f"No tasks found for worker '{worker_id}'") + + def main(): app() diff --git a/hindsight-api/hindsight_api/alembic/versions/l7g8h9i0j1k2_add_worker_columns.py b/hindsight-api/hindsight_api/alembic/versions/l7g8h9i0j1k2_add_worker_columns.py new file mode 100644 index 00000000..db60510d --- /dev/null +++ b/hindsight-api/hindsight_api/alembic/versions/l7g8h9i0j1k2_add_worker_columns.py @@ -0,0 +1,109 @@ +"""add_worker_columns + +Revision ID: l7g8h9i0j1k2 +Revises: k6f7g8h9i0j1 +Create Date: 2026-01-19 00:00:00.000000 + +This migration adds columns to async_operations for distributed worker support: +- worker_id: ID of the worker that claimed the task +- claimed_at: When the task was claimed +- retry_count: Number of retry attempts +- task_payload: The serialized task dictionary +""" + +from collections.abc import Sequence + +import sqlalchemy as sa +from alembic import context, op +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic. +revision: str = "l7g8h9i0j1k2" +down_revision: str | Sequence[str] | None = "k6f7g8h9i0j1" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def _get_schema_prefix() -> str: + """Get schema prefix for table names (required for multi-tenant support).""" + schema = context.config.get_main_option("target_schema") + return f'"{schema}".' if schema else "" + + +def upgrade() -> None: + """Add worker columns to async_operations.""" + schema = _get_schema_prefix() + + # Add worker_id column (ID of worker that claimed the task) + op.add_column( + "async_operations", + sa.Column("worker_id", sa.Text(), nullable=True), + schema=context.config.get_main_option("target_schema") or None, + ) + + # Add claimed_at column (when task was claimed by worker) + op.add_column( + "async_operations", + sa.Column("claimed_at", postgresql.TIMESTAMP(timezone=True), nullable=True), + schema=context.config.get_main_option("target_schema") or None, + ) + + # Add retry_count column (number of retry attempts) + op.add_column( + "async_operations", + sa.Column("retry_count", sa.Integer(), server_default="0", nullable=False), + schema=context.config.get_main_option("target_schema") or None, + ) + + # Add task_payload column (serialized task dictionary) + op.add_column( + "async_operations", + sa.Column( + "task_payload", + postgresql.JSONB(astext_type=sa.Text()), + nullable=True, + ), + schema=context.config.get_main_option("target_schema") or None, + ) + + # Add index for efficient worker polling (pending tasks ordered by creation time) + op.execute( + f"CREATE INDEX idx_async_operations_pending_claim ON {schema}async_operations (status, created_at) " + f"WHERE status = 'pending' AND task_payload IS NOT NULL" + ) + + # Add index for finding tasks by worker_id (for decommissioning) + op.execute( + f"CREATE INDEX idx_async_operations_worker_id ON {schema}async_operations (worker_id) WHERE worker_id IS NOT NULL" + ) + + +def downgrade() -> None: + """Remove worker columns from async_operations.""" + schema = _get_schema_prefix() + + # Drop indexes + op.execute(f"DROP INDEX IF EXISTS {schema}idx_async_operations_pending_claim") + op.execute(f"DROP INDEX IF EXISTS {schema}idx_async_operations_worker_id") + + # Drop columns + op.drop_column( + "async_operations", + "task_payload", + schema=context.config.get_main_option("target_schema") or None, + ) + op.drop_column( + "async_operations", + "retry_count", + schema=context.config.get_main_option("target_schema") or None, + ) + op.drop_column( + "async_operations", + "claimed_at", + schema=context.config.get_main_option("target_schema") or None, + ) + op.drop_column( + "async_operations", + "worker_id", + schema=context.config.get_main_option("target_schema") or None, + ) diff --git a/hindsight-api/hindsight_api/api/http.py b/hindsight-api/hindsight_api/api/http.py index 1e093b50..e7c3ab6c 100644 --- a/hindsight-api/hindsight_api/api/http.py +++ b/hindsight-api/hindsight_api/api/http.py @@ -1408,6 +1408,16 @@ def create_app( Lifespan context manager for startup and shutdown events. Note: This only fires when running the app standalone, not when mounted. """ + import asyncio + import socket + + from hindsight_api.config import get_config + from hindsight_api.worker import WorkerPoller + + config = get_config() + poller = None + poller_task = None + # Initialize OpenTelemetry metrics try: prometheus_reader = initialize_metrics(service_name="hindsight-api", service_version="1.0.0") @@ -1430,6 +1440,20 @@ def create_app( metrics_collector.set_db_pool(memory._pool) logging.info("DB pool metrics configured") + # Start worker poller if enabled (standalone mode) + if config.worker_enabled and memory._pool is not None: + worker_id = config.worker_id or socket.gethostname() + poller = WorkerPoller( + pool=memory._pool, + worker_id=worker_id, + executor=memory.execute_task, + poll_interval_ms=config.worker_poll_interval_ms, + batch_size=config.worker_batch_size, + max_retries=config.worker_max_retries, + ) + poller_task = asyncio.create_task(poller.run()) + logging.info(f"Worker poller started (worker_id={worker_id})") + # Call HTTP extension startup hook if http_extension: await http_extension.on_startup() @@ -1437,6 +1461,17 @@ def create_app( yield + # Shutdown worker poller if running + if poller is not None: + await poller.shutdown_graceful(timeout=30.0) + if poller_task is not None: + poller_task.cancel() + try: + await poller_task + except asyncio.CancelledError: + pass + logging.info("Worker poller stopped") + # Call HTTP extension shutdown hook if http_extension: await http_extension.on_shutdown() diff --git a/hindsight-api/hindsight_api/config.py b/hindsight-api/hindsight_api/config.py index 69a99f40..d0992812 100644 --- a/hindsight-api/hindsight_api/config.py +++ b/hindsight-api/hindsight_api/config.py @@ -106,10 +106,13 @@ ENV_DB_POOL_MAX_SIZE = "HINDSIGHT_API_DB_POOL_MAX_SIZE" ENV_DB_COMMAND_TIMEOUT = "HINDSIGHT_API_DB_COMMAND_TIMEOUT" ENV_DB_ACQUIRE_TIMEOUT = "HINDSIGHT_API_DB_ACQUIRE_TIMEOUT" -# Background task processing -ENV_TASK_BACKEND = "HINDSIGHT_API_TASK_BACKEND" -ENV_TASK_BACKEND_MEMORY_BATCH_SIZE = "HINDSIGHT_API_TASK_BACKEND_MEMORY_BATCH_SIZE" -ENV_TASK_BACKEND_MEMORY_BATCH_INTERVAL = "HINDSIGHT_API_TASK_BACKEND_MEMORY_BATCH_INTERVAL" +# Worker configuration (distributed task processing) +ENV_WORKER_ENABLED = "HINDSIGHT_API_WORKER_ENABLED" +ENV_WORKER_ID = "HINDSIGHT_API_WORKER_ID" +ENV_WORKER_POLL_INTERVAL_MS = "HINDSIGHT_API_WORKER_POLL_INTERVAL_MS" +ENV_WORKER_MAX_RETRIES = "HINDSIGHT_API_WORKER_MAX_RETRIES" +ENV_WORKER_BATCH_SIZE = "HINDSIGHT_API_WORKER_BATCH_SIZE" +ENV_WORKER_HTTP_PORT = "HINDSIGHT_API_WORKER_HTTP_PORT" # Reflect agent settings ENV_REFLECT_MAX_ITERATIONS = "HINDSIGHT_API_REFLECT_MAX_ITERATIONS" @@ -177,10 +180,13 @@ DEFAULT_DB_POOL_MAX_SIZE = 100 DEFAULT_DB_COMMAND_TIMEOUT = 60 # seconds DEFAULT_DB_ACQUIRE_TIMEOUT = 30 # seconds -# Background task processing -DEFAULT_TASK_BACKEND = "memory" # Options: "memory", "noop" -DEFAULT_TASK_BACKEND_MEMORY_BATCH_SIZE = 10 -DEFAULT_TASK_BACKEND_MEMORY_BATCH_INTERVAL = 1.0 # seconds +# Worker configuration (distributed task processing) +DEFAULT_WORKER_ENABLED = True # API runs worker by default (standalone mode) +DEFAULT_WORKER_ID = None # Will use hostname if not specified +DEFAULT_WORKER_POLL_INTERVAL_MS = 500 # Poll database every 500ms +DEFAULT_WORKER_MAX_RETRIES = 3 # Max retries before marking task failed +DEFAULT_WORKER_BATCH_SIZE = 10 # Tasks to claim per poll cycle +DEFAULT_WORKER_HTTP_PORT = 8889 # HTTP port for worker metrics/health # Reflect agent settings DEFAULT_REFLECT_MAX_ITERATIONS = 10 # Max tool call iterations before forcing response @@ -331,10 +337,13 @@ class HindsightConfig: db_command_timeout: int db_acquire_timeout: int - # Background task processing - task_backend: str - task_backend_memory_batch_size: int - task_backend_memory_batch_interval: float + # Worker configuration (distributed task processing) + worker_enabled: bool + worker_id: str | None + worker_poll_interval_ms: int + worker_max_retries: int + worker_batch_size: int + worker_http_port: int # Reflect agent settings reflect_max_iterations: int @@ -424,14 +433,13 @@ class HindsightConfig: db_pool_max_size=int(os.getenv(ENV_DB_POOL_MAX_SIZE, str(DEFAULT_DB_POOL_MAX_SIZE))), db_command_timeout=int(os.getenv(ENV_DB_COMMAND_TIMEOUT, str(DEFAULT_DB_COMMAND_TIMEOUT))), db_acquire_timeout=int(os.getenv(ENV_DB_ACQUIRE_TIMEOUT, str(DEFAULT_DB_ACQUIRE_TIMEOUT))), - # Background task processing - task_backend=os.getenv(ENV_TASK_BACKEND, DEFAULT_TASK_BACKEND), - task_backend_memory_batch_size=int( - os.getenv(ENV_TASK_BACKEND_MEMORY_BATCH_SIZE, str(DEFAULT_TASK_BACKEND_MEMORY_BATCH_SIZE)) - ), - task_backend_memory_batch_interval=float( - os.getenv(ENV_TASK_BACKEND_MEMORY_BATCH_INTERVAL, str(DEFAULT_TASK_BACKEND_MEMORY_BATCH_INTERVAL)) - ), + # Worker configuration + worker_enabled=os.getenv(ENV_WORKER_ENABLED, str(DEFAULT_WORKER_ENABLED)).lower() == "true", + worker_id=os.getenv(ENV_WORKER_ID) or DEFAULT_WORKER_ID, + worker_poll_interval_ms=int(os.getenv(ENV_WORKER_POLL_INTERVAL_MS, str(DEFAULT_WORKER_POLL_INTERVAL_MS))), + worker_max_retries=int(os.getenv(ENV_WORKER_MAX_RETRIES, str(DEFAULT_WORKER_MAX_RETRIES))), + worker_batch_size=int(os.getenv(ENV_WORKER_BATCH_SIZE, str(DEFAULT_WORKER_BATCH_SIZE))), + worker_http_port=int(os.getenv(ENV_WORKER_HTTP_PORT, str(DEFAULT_WORKER_HTTP_PORT))), # Reflect agent settings reflect_max_iterations=int(os.getenv(ENV_REFLECT_MAX_ITERATIONS, str(DEFAULT_REFLECT_MAX_ITERATIONS))), ) diff --git a/hindsight-api/hindsight_api/engine/memory_engine.py b/hindsight-api/hindsight_api/engine/memory_engine.py index 5d846374..8ddb96ac 100644 --- a/hindsight-api/hindsight_api/engine/memory_engine.py +++ b/hindsight-api/hindsight_api/engine/memory_engine.py @@ -159,7 +159,7 @@ from .retain.types import RetainContentDict from .search import think_utils from .search.reranking import CrossEncoderReranker from .search.tags import TagsMatch -from .task_backend import AsyncIOQueueBackend, NoopTaskBackend, TaskBackend +from .task_backend import BrokerTaskBackend, SyncTaskBackend, TaskBackend class Budget(str, Enum): @@ -229,8 +229,6 @@ class MemoryEngine(MemoryEngineInterface): db_command_timeout: int | None = None, db_acquire_timeout: int | None = None, task_backend: TaskBackend | None = None, - task_batch_size: int | None = None, - task_batch_interval: float | None = None, run_migrations: bool = True, operation_validator: "OperationValidatorExtension | None" = None, tenant_extension: "TenantExtension | None" = None, @@ -265,9 +263,7 @@ class MemoryEngine(MemoryEngineInterface): pool_max_size: Maximum number of connections in the pool. Defaults to HINDSIGHT_API_DB_POOL_MAX_SIZE. db_command_timeout: PostgreSQL command timeout in seconds. Defaults to HINDSIGHT_API_DB_COMMAND_TIMEOUT. db_acquire_timeout: Connection acquisition timeout in seconds. Defaults to HINDSIGHT_API_DB_ACQUIRE_TIMEOUT. - task_backend: Custom task backend. If not provided, uses AsyncIOQueueBackend. - task_batch_size: Background task batch size. Defaults to HINDSIGHT_API_TASK_BATCH_SIZE. - task_batch_interval: Background task batch interval in seconds. Defaults to HINDSIGHT_API_TASK_BATCH_INTERVAL. + task_backend: Custom task backend. If not provided, uses BrokerTaskBackend for distributed processing. run_migrations: Whether to run database migrations during initialize(). Default: True operation_validator: Optional extension to validate operations before execution. If provided, retain/recall/reflect operations will be validated. @@ -405,13 +401,9 @@ class MemoryEngine(MemoryEngineInterface): self._cross_encoder_reranker = CrossEncoderReranker(cross_encoder=cross_encoder) # Initialize task backend - _task_batch_size = task_batch_size if task_batch_size is not None else config.task_backend_memory_batch_size - _task_batch_interval = ( - task_batch_interval if task_batch_interval is not None else config.task_backend_memory_batch_interval - ) - self._task_backend = task_backend or AsyncIOQueueBackend( - batch_size=_task_batch_size, batch_interval=_task_batch_interval - ) + # If no custom backend provided, use BrokerTaskBackend which stores tasks in PostgreSQL + # The pool_getter lambda will return the pool once it's initialized + self._task_backend = task_backend or BrokerTaskBackend(pool_getter=lambda: self._pool) # Backpressure mechanism: limit concurrent searches to prevent overwhelming the database # Configurable via HINDSIGHT_API_RECALL_MAX_CONCURRENT (default: 50) diff --git a/hindsight-api/hindsight_api/engine/task_backend.py b/hindsight-api/hindsight_api/engine/task_backend.py index 1e071958..29420463 100644 --- a/hindsight-api/hindsight_api/engine/task_backend.py +++ b/hindsight-api/hindsight_api/engine/task_backend.py @@ -1,31 +1,40 @@ """ -Abstract task backend for running async tasks. +Task backend for distributed task processing. -This provides an abstraction that can be adapted to different execution models: -- AsyncIO queue (default implementation) -- Pub/Sub architectures (future) -- Message brokers (future) +This provides an abstraction for task storage and execution: +- BrokerTaskBackend: Uses PostgreSQL as broker (production) +- SyncTaskBackend: Executes tasks immediately (testing/embedded) """ -import asyncio +import json import logging from abc import ABC, abstractmethod from collections.abc import Awaitable, Callable -from typing import Any +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + import asyncpg logger = logging.getLogger(__name__) +def fq_table(table: str, schema: str | None = None) -> str: + """Get fully-qualified table name with optional schema prefix.""" + if schema: + return f'"{schema}".{table}' + return table + + class TaskBackend(ABC): """ Abstract base class for task execution backends. Implementations must: 1. Store/publish task events (as serializable dicts) - 2. Execute tasks through a provided executor callback + 2. Execute tasks through a provided executor callback (optional) The backend treats tasks as pure dictionaries that can be serialized - and sent over the network. The executor (typically MemoryEngine.execute_task) + and stored in the database. The executor (typically MemoryEngine.execute_task) receives the dict and routes it to the appropriate handler. """ @@ -46,7 +55,7 @@ class TaskBackend(ABC): @abstractmethod async def initialize(self): """ - Initialize the backend (e.g., start workers, connect to broker). + Initialize the backend (e.g., connect to database). """ pass @@ -63,7 +72,7 @@ class TaskBackend(ABC): @abstractmethod async def shutdown(self): """ - Shutdown the backend gracefully (e.g., stop workers, close connections). + Shutdown the backend gracefully. """ pass @@ -93,9 +102,8 @@ class SyncTaskBackend(TaskBackend): """ Synchronous task backend that executes tasks immediately. - This is useful for embedded/CLI usage where we don't want background - workers that prevent clean exit. Tasks are executed inline rather than - being queued. + This is useful for tests and embedded/CLI usage where we don't want + background workers. Tasks are executed inline rather than being queued. """ async def initialize(self): @@ -121,221 +129,123 @@ class SyncTaskBackend(TaskBackend): logger.debug("SyncTaskBackend shutdown") -class NoopTaskBackend(TaskBackend): +class BrokerTaskBackend(TaskBackend): """ - No-op task backend that discards all tasks. + Task backend using PostgreSQL as broker. - This is useful for tests where background task execution is not needed - and would only slow down the test suite. + submit_task() stores task_payload in async_operations table. + Actual polling and execution is handled separately by WorkerPoller. + + This backend is used by the API to store tasks. Workers poll + the database separately to claim and execute tasks. """ - async def initialize(self): - """No-op.""" - self._initialized = True - logger.debug("NoopTaskBackend initialized") - - async def submit_task(self, task_dict: dict[str, Any]): - """Discard the task (do nothing).""" - pass - - async def shutdown(self): - """No-op.""" - self._initialized = False - logger.debug("NoopTaskBackend shutdown") - - -class AsyncIOQueueBackend(TaskBackend): - """ - Task backend implementation using asyncio queues. - - This is the default implementation that uses in-process asyncio queues - and a periodic consumer worker. - """ - - def __init__(self, batch_size: int = 10, batch_interval: float = 1.0): + def __init__( + self, + pool_getter: Callable[[], "asyncpg.Pool"], + schema: str | None = None, + ): """ - Initialize AsyncIO queue backend. + Initialize the broker task backend. Args: - batch_size: Maximum number of tasks to process in one batch - batch_interval: Maximum time (seconds) to wait before processing batch + pool_getter: Callable that returns the asyncpg connection pool + schema: Database schema for multi-tenant support (optional) """ super().__init__() - self._queue: asyncio.Queue | None = None - self._worker_task: asyncio.Task | None = None - self._shutdown_event: asyncio.Event | None = None - self._batch_size = batch_size - self._batch_interval = batch_interval - self._in_flight_count = 0 - self._in_flight_lock = asyncio.Lock() + self._pool_getter = pool_getter + self._schema = schema async def initialize(self): - """Initialize the queue and start the worker.""" - if self._initialized: - return - - self._queue = asyncio.Queue() - self._shutdown_event = asyncio.Event() - self._worker_task = asyncio.create_task(self._worker()) + """Initialize the backend.""" self._initialized = True - logger.info("AsyncIOQueueBackend initialized") + logger.info("BrokerTaskBackend initialized") async def submit_task(self, task_dict: dict[str, Any]): """ - Submit a task by putting it in the queue. + Store task payload in async_operations table. + + The task_dict should contain an 'operation_id' if updating an existing + operation record, otherwise a new operation will be created. Args: - task_dict: Task dictionary to execute + task_dict: Task dictionary to store (must be JSON serializable) """ if not self._initialized: await self.initialize() - await self._queue.put(task_dict) + pool = self._pool_getter() + operation_id = task_dict.get("operation_id") + task_type = task_dict.get("type", "unknown") + bank_id = task_dict.get("bank_id") + payload_json = json.dumps(task_dict) + + table = fq_table("async_operations", self._schema) + + if operation_id: + # Update existing operation with task payload + await pool.execute( + f""" + UPDATE {table} + SET task_payload = $1::jsonb, updated_at = now() + WHERE operation_id = $2 + """, + payload_json, + operation_id, + ) + logger.debug(f"Updated task payload for operation {operation_id}") + else: + # Insert new operation (for tasks without pre-created records) + # e.g., access_count_update tasks + import uuid + + new_id = uuid.uuid4() + await pool.execute( + f""" + INSERT INTO {table} (operation_id, bank_id, operation_type, status, task_payload) + VALUES ($1, $2, $3, 'pending', $4::jsonb) + """, + new_id, + bank_id, + task_type, + payload_json, + ) + logger.debug(f"Created new operation {new_id} for task type {task_type}") + + async def shutdown(self): + """Shutdown the backend.""" + self._initialized = False + logger.info("BrokerTaskBackend shutdown") async def wait_for_pending_tasks(self, timeout: float = 120.0): """ - Wait for all pending tasks in the queue and in-flight tasks to complete. + Wait for pending tasks to be processed. - This is useful in tests to ensure background tasks complete before assertions. + In the broker model, this polls the database to check if tasks + for this process have been completed. This is useful in tests + when worker_enabled=True (API processes its own tasks). Args: - timeout: Maximum time to wait in seconds (default 120s for long-running tasks) + timeout: Maximum time to wait in seconds """ - if not self._initialized or self._queue is None: - return + import asyncio + + pool = self._pool_getter() + table = fq_table("async_operations", self._schema) - # Wait for queue to be empty AND no in-flight tasks start_time = asyncio.get_event_loop().time() while asyncio.get_event_loop().time() - start_time < timeout: - async with self._in_flight_lock: - in_flight = self._in_flight_count + # Check if there are any pending tasks with payloads + count = await pool.fetchval( + f""" + SELECT COUNT(*) FROM {table} + WHERE status = 'pending' AND task_payload IS NOT NULL + """ + ) - if self._queue.empty() and in_flight == 0: - # Queue is empty and no tasks in flight, we're done + if count == 0: return - # Wait a bit before checking again await asyncio.sleep(0.5) - async def shutdown(self): - """Shutdown the worker and drain the queue.""" - if not self._initialized: - return - - logger.info("Shutting down AsyncIOQueueBackend...") - - # Signal shutdown - self._shutdown_event.set() - - # Cancel worker - if self._worker_task is not None: - self._worker_task.cancel() - try: - await self._worker_task - except asyncio.CancelledError: - pass # Worker cancelled successfully - - self._initialized = False - logger.info("AsyncIOQueueBackend shutdown complete") - - async def _execute_task_with_tracking(self, task_dict: dict[str, Any]): - """Execute a task and track its in-flight status.""" - async with self._in_flight_lock: - self._in_flight_count += 1 - try: - await self._execute_task(task_dict) - finally: - async with self._in_flight_lock: - self._in_flight_count -= 1 - - async def _execute_task_no_tracking(self, task_dict: dict[str, Any]): - """Execute a task without in-flight tracking (tracking done at batch level).""" - await self._execute_task(task_dict) - - def _get_queue_stats(self) -> tuple[int, dict[str, int]]: - """Get current queue size and bank_id distribution.""" - queue_size = self._queue.qsize() if self._queue else 0 - bank_distribution: dict[str, int] = {} - - if queue_size > 0 and self._queue: - # Peek at queue items without removing them - # Note: This is a snapshot and may not be perfectly accurate due to concurrency - try: - # Access internal deque for logging purposes only - items = list(self._queue._queue) # type: ignore[attr-defined] - for item in items: - bank_id = item.get("bank_id", "unknown") - bank_distribution[bank_id] = bank_distribution.get(bank_id, 0) + 1 - except Exception: - pass # Queue access failed, return empty distribution - - return queue_size, bank_distribution - - async def _worker(self): - """ - Background worker that processes tasks in batches. - - Collects tasks for up to batch_interval seconds or batch_size items, - then processes them. - """ - while not self._shutdown_event.is_set(): - try: - # Collect tasks for batching - tasks = [] - deadline = asyncio.get_event_loop().time() + self._batch_interval - - while len(tasks) < self._batch_size and asyncio.get_event_loop().time() < deadline: - try: - remaining_time = max(0.1, deadline - asyncio.get_event_loop().time()) - task_dict = await asyncio.wait_for(self._queue.get(), timeout=remaining_time) - # Track task as in-flight immediately when picked up from queue - # This prevents wait_for_pending_tasks from returning too early - async with self._in_flight_lock: - self._in_flight_count += 1 - tasks.append(task_dict) - except TimeoutError: - break - - # Process batch - if tasks: - # Log batch start with queue stats - queue_size, bank_distribution = self._get_queue_stats() - - # Summarize batch by task type and bank - batch_summary: dict[str, dict[str, int]] = {} - for task_dict in tasks: - task_type = task_dict.get("type", "unknown") - bank_id = task_dict.get("bank_id", "unknown") - if task_type not in batch_summary: - batch_summary[task_type] = {} - batch_summary[task_type][bank_id] = batch_summary[task_type].get(bank_id, 0) + 1 - - # Build log message - batch_parts = [] - for task_type, banks in sorted(batch_summary.items()): - bank_str = ", ".join(f"{b}:{c}" for b, c in sorted(banks.items())) - batch_parts.append(f"{task_type}[{bank_str}]") - batch_str = ", ".join(batch_parts) - - if queue_size > 0: - pending_str = ", ".join(f"{k}:{v}" for k, v in sorted(bank_distribution.items())) - logger.info( - f"Processing {len(tasks)} tasks: {batch_str} (pending={queue_size} [{pending_str}])" - ) - else: - logger.info(f"Processing {len(tasks)} tasks: {batch_str}") - - # Execute tasks concurrently (in_flight already tracked when picked up) - await asyncio.gather( - *[self._execute_task_no_tracking(task_dict) for task_dict in tasks], return_exceptions=True - ) - - # Decrement in_flight count after all tasks complete - async with self._in_flight_lock: - self._in_flight_count -= len(tasks) - - except asyncio.CancelledError: - break - except Exception as e: - logger.error(f"Worker error: {e}") - await asyncio.sleep(1) # Backoff on error + logger.warning(f"Timeout waiting for pending tasks after {timeout}s") diff --git a/hindsight-api/hindsight_api/main.py b/hindsight-api/hindsight_api/main.py index 95bc613c..44a2872b 100644 --- a/hindsight-api/hindsight_api/main.py +++ b/hindsight-api/hindsight_api/main.py @@ -219,9 +219,12 @@ def main(): db_pool_max_size=config.db_pool_max_size, db_command_timeout=config.db_command_timeout, db_acquire_timeout=config.db_acquire_timeout, - task_backend=config.task_backend, - task_backend_memory_batch_size=config.task_backend_memory_batch_size, - task_backend_memory_batch_interval=config.task_backend_memory_batch_interval, + worker_enabled=config.worker_enabled, + worker_id=config.worker_id, + worker_poll_interval_ms=config.worker_poll_interval_ms, + worker_max_retries=config.worker_max_retries, + worker_batch_size=config.worker_batch_size, + worker_http_port=config.worker_http_port, reflect_max_iterations=config.reflect_max_iterations, mental_model_refresh_concurrency=config.mental_model_refresh_concurrency, ) diff --git a/hindsight-api/hindsight_api/worker/__init__.py b/hindsight-api/hindsight_api/worker/__init__.py new file mode 100644 index 00000000..8eb27f9f --- /dev/null +++ b/hindsight-api/hindsight_api/worker/__init__.py @@ -0,0 +1,11 @@ +""" +Worker package for distributed task processing. + +This package provides: +- WorkerPoller: Polls PostgreSQL for pending tasks and executes them +- main: CLI entry point for hindsight-worker +""" + +from .poller import WorkerPoller + +__all__ = ["WorkerPoller"] diff --git a/hindsight-api/hindsight_api/worker/main.py b/hindsight-api/hindsight_api/worker/main.py new file mode 100644 index 00000000..a230b73e --- /dev/null +++ b/hindsight-api/hindsight_api/worker/main.py @@ -0,0 +1,285 @@ +""" +Command-line interface for Hindsight Worker. + +Run the worker with: + hindsight-worker + +Stop with Ctrl+C (graceful shutdown). +""" + +import argparse +import asyncio +import atexit +import logging +import os +import signal +import socket +import sys +import warnings + +from ..config import get_config +from ..engine.task_backend import SyncTaskBackend +from .poller import WorkerPoller + +# Filter deprecation warnings from third-party libraries +warnings.filterwarnings("ignore", message="websockets.legacy is deprecated") +warnings.filterwarnings("ignore", message="websockets.server.WebSocketServerProtocol is deprecated") + +# Disable tokenizers parallelism to avoid warnings +os.environ["TOKENIZERS_PARALLELISM"] = "false" + +logger = logging.getLogger(__name__) + + +def create_worker_app(poller: WorkerPoller, memory): + """Create a minimal FastAPI app for worker metrics and health.""" + from fastapi import FastAPI + from fastapi.responses import JSONResponse, Response + from prometheus_client import CONTENT_TYPE_LATEST, generate_latest + + from ..metrics import create_metrics_collector, get_metrics_collector, initialize_metrics + + app = FastAPI( + title="Hindsight Worker", + description="Worker process for distributed task execution", + ) + + # Initialize OpenTelemetry metrics + try: + prometheus_reader = initialize_metrics(service_name="hindsight-worker", service_version="1.0.0") + create_metrics_collector() + app.state.prometheus_reader = prometheus_reader + logger.info("Metrics initialized - available at /metrics endpoint") + except Exception as e: + logger.warning(f"Failed to initialize metrics: {e}. Metrics will be disabled.") + app.state.prometheus_reader = None + + # Set up DB pool metrics if available + metrics_collector = get_metrics_collector() + if memory._pool is not None and hasattr(metrics_collector, "set_db_pool"): + metrics_collector.set_db_pool(memory._pool) + logger.info("DB pool metrics configured") + + @app.get( + "/health", + summary="Health check endpoint", + description="Returns worker health status including database connectivity", + tags=["Monitoring"], + ) + async def health_endpoint(): + """Health check endpoint.""" + health = await memory.health_check() + health["worker_id"] = poller.worker_id + health["is_shutdown"] = poller.is_shutdown + status_code = 200 if health.get("status") == "healthy" else 503 + return JSONResponse(content=health, status_code=status_code) + + @app.get( + "/metrics", + summary="Prometheus metrics endpoint", + description="Exports metrics in Prometheus format for scraping", + tags=["Monitoring"], + ) + async def metrics_endpoint(): + """Return Prometheus metrics.""" + metrics_data = generate_latest() + return Response(content=metrics_data, media_type=CONTENT_TYPE_LATEST) + + @app.get( + "/", + summary="Worker info", + description="Basic worker information", + tags=["Info"], + ) + async def root(): + """Return basic worker info.""" + return { + "service": "hindsight-worker", + "worker_id": poller.worker_id, + "is_shutdown": poller.is_shutdown, + } + + return app + + +def main(): + """Main entry point for the hindsight-worker CLI.""" + # Load configuration from environment + config = get_config() + + parser = argparse.ArgumentParser( + prog="hindsight-worker", + description="Hindsight Worker - distributed task processor", + ) + + # Worker options + parser.add_argument( + "--worker-id", + default=config.worker_id or socket.gethostname(), + help="Worker identifier (default: hostname, env: HINDSIGHT_API_WORKER_ID)", + ) + parser.add_argument( + "--poll-interval", + type=int, + default=config.worker_poll_interval_ms, + help=f"Poll interval in milliseconds (default: {config.worker_poll_interval_ms}, env: HINDSIGHT_API_WORKER_POLL_INTERVAL_MS)", + ) + parser.add_argument( + "--batch-size", + type=int, + default=config.worker_batch_size, + help=f"Tasks to claim per poll (default: {config.worker_batch_size}, env: HINDSIGHT_API_WORKER_BATCH_SIZE)", + ) + parser.add_argument( + "--max-retries", + type=int, + default=config.worker_max_retries, + help=f"Max retries before marking failed (default: {config.worker_max_retries}, env: HINDSIGHT_API_WORKER_MAX_RETRIES)", + ) + + # HTTP server options + parser.add_argument( + "--http-port", + type=int, + default=config.worker_http_port, + help=f"HTTP port for metrics/health endpoints (default: {config.worker_http_port}, env: HINDSIGHT_API_WORKER_HTTP_PORT)", + ) + parser.add_argument( + "--http-host", + default="0.0.0.0", + help="HTTP host to bind (default: 0.0.0.0)", + ) + + # Logging options + parser.add_argument( + "--log-level", + default=config.log_level, + choices=["critical", "error", "warning", "info", "debug", "trace"], + help=f"Log level (default: {config.log_level}, env: HINDSIGHT_API_LOG_LEVEL)", + ) + + args = parser.parse_args() + + # Configure logging + config.configure_logging() + + # Import MemoryEngine here to avoid circular imports + from .. import MemoryEngine + + print(f"Starting Hindsight Worker: {args.worker_id}") + print(f" Poll interval: {args.poll_interval}ms") + print(f" Batch size: {args.batch_size}") + print(f" Max retries: {args.max_retries}") + print(f" HTTP server: {args.http_host}:{args.http_port}") + print() + + # Global references for cleanup + memory = None + poller = None + + async def run(): + nonlocal memory, poller + import uvicorn + + # Initialize MemoryEngine + # Workers use SyncTaskBackend because they execute tasks directly, + # they don't need to store tasks (they poll from DB) + memory = MemoryEngine( + run_migrations=False, # Workers don't run migrations + task_backend=SyncTaskBackend(), + ) + + await memory.initialize() + + print(f"Database connected: {config.database_url}") + + # Create and start the poller + poller = WorkerPoller( + pool=memory._pool, + worker_id=args.worker_id, + executor=memory.execute_task, + poll_interval_ms=args.poll_interval, + batch_size=args.batch_size, + max_retries=args.max_retries, + ) + + # Create the HTTP app for metrics/health + app = create_worker_app(poller, memory) + + # Setup signal handlers for graceful shutdown + shutdown_requested = asyncio.Event() + + def signal_handler(signum, frame): + print(f"\nReceived signal {signum}, initiating graceful shutdown...") + shutdown_requested.set() + + signal.signal(signal.SIGINT, signal_handler) + signal.signal(signal.SIGTERM, signal_handler) + + # Create uvicorn config and server + uvicorn_config = uvicorn.Config( + app, + host=args.http_host, + port=args.http_port, + log_level="info", # Reduce uvicorn noise + access_log=False, + ) + server = uvicorn.Server(uvicorn_config) + + # Run the poller and HTTP server concurrently + poller_task = asyncio.create_task(poller.run()) + http_task = asyncio.create_task(server.serve()) + + print(f"Worker started. Metrics available at http://{args.http_host}:{args.http_port}/metrics") + + # Wait for shutdown signal + await shutdown_requested.wait() + + # Graceful shutdown + print("Shutting down HTTP server...") + server.should_exit = True + + print("Waiting for poller to finish...") + await poller.shutdown_graceful(timeout=30.0) + poller_task.cancel() + try: + await poller_task + except asyncio.CancelledError: + pass + + # Wait for HTTP server to finish + try: + await asyncio.wait_for(http_task, timeout=5.0) + except asyncio.TimeoutError: + http_task.cancel() + try: + await http_task + except asyncio.CancelledError: + pass + + # Close memory engine + await memory.close() + print("Worker shutdown complete") + + def cleanup(): + """Synchronous cleanup for atexit.""" + if memory is not None and memory._pg0 is not None: + try: + loop = asyncio.new_event_loop() + loop.run_until_complete(memory._pg0.stop()) + loop.close() + print("\npg0 stopped.") + except Exception as e: + print(f"\nError stopping pg0: {e}") + + atexit.register(cleanup) + + try: + asyncio.run(run()) + except KeyboardInterrupt: + print("\nWorker interrupted") + sys.exit(0) + + +if __name__ == "__main__": + main() diff --git a/hindsight-api/hindsight_api/worker/poller.py b/hindsight-api/hindsight_api/worker/poller.py new file mode 100644 index 00000000..e85f8a18 --- /dev/null +++ b/hindsight-api/hindsight_api/worker/poller.py @@ -0,0 +1,281 @@ +""" +Worker poller for distributed task execution. + +Polls PostgreSQL for pending tasks and executes them using +FOR UPDATE SKIP LOCKED for safe concurrent claiming. +""" + +import asyncio +import json +import logging +import traceback +from collections.abc import Awaitable, Callable +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + import asyncpg + +logger = logging.getLogger(__name__) + + +def fq_table(table: str, schema: str | None = None) -> str: + """Get fully-qualified table name with optional schema prefix.""" + if schema: + return f'"{schema}".{table}' + return table + + +class WorkerPoller: + """ + Polls PostgreSQL for pending tasks and executes them. + + Uses FOR UPDATE SKIP LOCKED for safe distributed claiming, + allowing multiple workers to process tasks without conflicts. + """ + + def __init__( + self, + pool: "asyncpg.Pool", + worker_id: str, + executor: Callable[[dict[str, Any]], Awaitable[None]], + poll_interval_ms: int = 500, + batch_size: int = 10, + max_retries: int = 3, + schema: str | None = None, + ): + """ + Initialize the worker poller. + + Args: + pool: asyncpg connection pool + worker_id: Unique identifier for this worker + executor: Async function to execute tasks (typically MemoryEngine.execute_task) + poll_interval_ms: Interval between polls when no tasks found (milliseconds) + batch_size: Maximum number of tasks to claim per poll cycle + max_retries: Maximum retry attempts before marking task as failed + schema: Database schema for multi-tenant support (optional) + """ + self._pool = pool + self._worker_id = worker_id + self._executor = executor + self._poll_interval_ms = poll_interval_ms + self._batch_size = batch_size + self._max_retries = max_retries + self._schema = schema + self._shutdown = asyncio.Event() + self._current_tasks: set[asyncio.Task] = set() + self._in_flight_count = 0 + self._in_flight_lock = asyncio.Lock() + + async def claim_batch(self) -> list[tuple[str, dict[str, Any]]]: + """ + Claim up to batch_size pending tasks atomically. + + Uses FOR UPDATE SKIP LOCKED to ensure no conflicts with other workers. + + Returns: + List of tuples (operation_id, task_dict) + """ + table = fq_table("async_operations", self._schema) + + async with self._pool.acquire() as conn: + async with conn.transaction(): + # Select and lock pending tasks + rows = await conn.fetch( + f""" + SELECT operation_id, task_payload + FROM {table} + WHERE status = 'pending' AND task_payload IS NOT NULL + ORDER BY created_at + LIMIT $1 + FOR UPDATE SKIP LOCKED + """, + self._batch_size, + ) + + if not rows: + return [] + + # Claim the tasks by updating status and worker_id + operation_ids = [row["operation_id"] for row in rows] + await conn.execute( + f""" + UPDATE {table} + SET status = 'processing', worker_id = $1, claimed_at = now(), updated_at = now() + WHERE operation_id = ANY($2) + """, + self._worker_id, + operation_ids, + ) + + # Parse and return task payloads + return [(str(row["operation_id"]), json.loads(row["task_payload"])) for row in rows] + + async def _mark_completed(self, operation_id: str): + """Mark a task as completed.""" + table = fq_table("async_operations", self._schema) + await self._pool.execute( + f""" + UPDATE {table} + SET status = 'completed', completed_at = now(), updated_at = now() + WHERE operation_id = $1 + """, + operation_id, + ) + + async def _mark_failed(self, operation_id: str, error_message: str): + """Mark a task as failed with error message.""" + table = fq_table("async_operations", self._schema) + # Truncate error message if too long (max 5000 chars in schema) + error_message = error_message[:5000] if len(error_message) > 5000 else error_message + await self._pool.execute( + f""" + UPDATE {table} + SET status = 'failed', error_message = $2, completed_at = now(), updated_at = now() + WHERE operation_id = $1 + """, + operation_id, + error_message, + ) + + async def _retry_or_fail(self, operation_id: str, error_message: str): + """Increment retry count or mark as failed if max retries exceeded.""" + table = fq_table("async_operations", self._schema) + + # Get current retry count + row = await self._pool.fetchrow( + f"SELECT retry_count FROM {table} WHERE operation_id = $1", + operation_id, + ) + + if row is None: + logger.warning(f"Operation {operation_id} not found, cannot retry") + return + + retry_count = row["retry_count"] + + if retry_count >= self._max_retries: + # Max retries exceeded, mark as failed + await self._mark_failed( + operation_id, f"Max retries ({self._max_retries}) exceeded. Last error: {error_message}" + ) + logger.error(f"Task {operation_id} failed after {retry_count} retries") + else: + # Increment retry and reset to pending + await self._pool.execute( + f""" + UPDATE {table} + SET status = 'pending', worker_id = NULL, claimed_at = NULL, + retry_count = retry_count + 1, updated_at = now() + WHERE operation_id = $1 + """, + operation_id, + ) + logger.warning(f"Task {operation_id} failed, will retry (attempt {retry_count + 1}/{self._max_retries})") + + async def execute_task(self, operation_id: str, task_dict: dict[str, Any]): + """Execute a single task and update its status.""" + task_type = task_dict.get("type", "unknown") + bank_id = task_dict.get("bank_id", "unknown") + + try: + logger.debug(f"Executing task {operation_id} (type={task_type}, bank={bank_id})") + await self._executor(task_dict) + await self._mark_completed(operation_id) + logger.debug(f"Task {operation_id} completed successfully") + except Exception as e: + error_msg = f"{type(e).__name__}: {e}\n{traceback.format_exc()}" + logger.error(f"Task {operation_id} failed: {e}") + await self._retry_or_fail(operation_id, error_msg) + + async def run(self): + """ + Main polling loop. + + Continuously polls for pending tasks, claims them, and executes them + until shutdown is signaled. + """ + logger.info(f"Worker {self._worker_id} starting polling loop") + + while not self._shutdown.is_set(): + try: + # Claim a batch of tasks + tasks = await self.claim_batch() + + if tasks: + # Log batch info + task_types = {} + for _, task_dict in tasks: + t = task_dict.get("type", "unknown") + task_types[t] = task_types.get(t, 0) + 1 + types_str = ", ".join(f"{k}:{v}" for k, v in task_types.items()) + logger.info(f"Worker {self._worker_id} claimed {len(tasks)} tasks: {types_str}") + + # Track in-flight tasks + async with self._in_flight_lock: + self._in_flight_count += len(tasks) + + # Execute tasks concurrently + try: + await asyncio.gather( + *[self.execute_task(op_id, task_dict) for op_id, task_dict in tasks], + return_exceptions=True, + ) + finally: + async with self._in_flight_lock: + self._in_flight_count -= len(tasks) + else: + # No tasks found, wait before polling again + try: + await asyncio.wait_for( + self._shutdown.wait(), + timeout=self._poll_interval_ms / 1000, + ) + except asyncio.TimeoutError: + pass # Normal timeout, continue polling + + except asyncio.CancelledError: + logger.info(f"Worker {self._worker_id} polling loop cancelled") + break + except Exception as e: + logger.error(f"Worker {self._worker_id} error in polling loop: {e}") + traceback.print_exc() + # Backoff on error + await asyncio.sleep(1) + + logger.info(f"Worker {self._worker_id} polling loop stopped") + + async def shutdown_graceful(self, timeout: float = 30.0): + """ + Signal shutdown and wait for current tasks to complete. + + Args: + timeout: Maximum time to wait for in-flight tasks (seconds) + """ + logger.info(f"Worker {self._worker_id} initiating graceful shutdown") + self._shutdown.set() + + # Wait for in-flight tasks to complete + start_time = asyncio.get_event_loop().time() + while asyncio.get_event_loop().time() - start_time < timeout: + async with self._in_flight_lock: + in_flight = self._in_flight_count + + if in_flight == 0: + logger.info(f"Worker {self._worker_id} graceful shutdown complete") + return + + logger.info(f"Worker {self._worker_id} waiting for {in_flight} in-flight tasks") + await asyncio.sleep(0.5) + + logger.warning(f"Worker {self._worker_id} shutdown timeout after {timeout}s") + + @property + def worker_id(self) -> str: + """Get the worker ID.""" + return self._worker_id + + @property + def is_shutdown(self) -> bool: + """Check if shutdown has been signaled.""" + return self._shutdown.is_set() diff --git a/hindsight-api/pyproject.toml b/hindsight-api/pyproject.toml index aaa7c355..f9f1fea8 100644 --- a/hindsight-api/pyproject.toml +++ b/hindsight-api/pyproject.toml @@ -63,6 +63,7 @@ test = [ [project.scripts] hindsight-api = "hindsight_api.main:main" +hindsight-worker = "hindsight_api.worker.main:main" hindsight-local-mcp = "hindsight_api.mcp_local:main" hindsight-admin = "hindsight_api.admin.cli:main" diff --git a/hindsight-api/tests/conftest.py b/hindsight-api/tests/conftest.py index d2db8dde..6fec5abc 100644 --- a/hindsight-api/tests/conftest.py +++ b/hindsight-api/tests/conftest.py @@ -12,6 +12,7 @@ from hindsight_api import MemoryEngine, LLMConfig, LocalSTEmbeddings, RequestCon from hindsight_api.engine.cross_encoder import LocalSTCrossEncoder from hindsight_api.engine.query_analyzer import DateparserQueryAnalyzer +from hindsight_api.engine.task_backend import SyncTaskBackend from hindsight_api.pg0 import EmbeddedPostgres # Default pg0 instance configuration for tests @@ -147,6 +148,7 @@ async def memory(pg0_db_url, embeddings, cross_encoder, query_analyzer): Uses pg0_db_url (a postgresql:// URL) directly, so MemoryEngine won't try to manage pg0 lifecycle - that's handled by the session-scoped pg0_db_url fixture. Migrations are disabled here since they're run once at session scope in pg0_db_url. + Uses SyncTaskBackend so async tasks execute immediately (no worker needed). """ mem = MemoryEngine( db_url=pg0_db_url, # Direct postgresql:// URL, not pg0:// @@ -160,6 +162,7 @@ async def memory(pg0_db_url, embeddings, cross_encoder, query_analyzer): pool_min_size=1, pool_max_size=5, run_migrations=False, # Migrations already run at session scope + task_backend=SyncTaskBackend(), # Execute tasks immediately in tests ) await mem.initialize() yield mem diff --git a/hindsight-api/tests/test_custom_embedding_dimension.py b/hindsight-api/tests/test_custom_embedding_dimension.py index fd572ac9..57b62c5e 100644 --- a/hindsight-api/tests/test_custom_embedding_dimension.py +++ b/hindsight-api/tests/test_custom_embedding_dimension.py @@ -17,6 +17,7 @@ from hindsight_api import MemoryEngine, RequestContext from hindsight_api.engine.embeddings import LocalSTEmbeddings, OpenAIEmbeddings, CohereEmbeddings from hindsight_api.engine.cross_encoder import LocalSTCrossEncoder, CohereCrossEncoder from hindsight_api.engine.query_analyzer import DateparserQueryAnalyzer +from hindsight_api.engine.task_backend import SyncTaskBackend from hindsight_api.extensions import TenantExtension, TenantContext from hindsight_api.migrations import run_migrations, ensure_embedding_dimension @@ -323,6 +324,7 @@ class TestOpenAIEmbeddings: pool_max_size=3, run_migrations=False, tenant_extension=SchemaTenantExtension(schema_name), + task_backend=SyncTaskBackend(), ) try: @@ -392,6 +394,7 @@ class TestOpenAIEmbeddings: pool_max_size=3, run_migrations=False, tenant_extension=SchemaTenantExtension(schema_name), + task_backend=SyncTaskBackend(), ) try: @@ -559,6 +562,7 @@ class TestCohereIntegration: pool_max_size=3, run_migrations=False, tenant_extension=SchemaTenantExtension(schema_name), + task_backend=SyncTaskBackend(), ) try: diff --git a/hindsight-api/tests/test_load_large_batch.py b/hindsight-api/tests/test_load_large_batch.py index c2f6d27f..9a413948 100644 --- a/hindsight-api/tests/test_load_large_batch.py +++ b/hindsight-api/tests/test_load_large_batch.py @@ -19,6 +19,7 @@ import pytest_asyncio from hindsight_api import MemoryEngine, LLMConfig, LocalSTEmbeddings, RequestContext from hindsight_api.engine.cross_encoder import LocalSTCrossEncoder from hindsight_api.engine.query_analyzer import DateparserQueryAnalyzer +from hindsight_api.engine.task_backend import SyncTaskBackend from hindsight_api.engine.retain.fact_extraction import FactExtractionResponse, ExtractedFact from hindsight_api.engine.llm_wrapper import TokenUsage @@ -106,6 +107,7 @@ class TestLargeBatchRetain: pool_max_size=10, run_migrations=False, skip_llm_verification=True, # Skip LLM verification since we're mocking + task_backend=SyncTaskBackend(), # Execute tasks immediately in tests ) await mem.initialize() yield mem diff --git a/hindsight-api/tests/test_worker.py b/hindsight-api/tests/test_worker.py new file mode 100644 index 00000000..f9748678 --- /dev/null +++ b/hindsight-api/tests/test_worker.py @@ -0,0 +1,592 @@ +""" +Tests for the distributed worker system. + +Tests cover: +- BrokerTaskBackend task submission and storage +- WorkerPoller task claiming with FOR UPDATE SKIP LOCKED +- Concurrent workers claiming different tasks (no duplicates) +- Task completion and failure handling +- Retry mechanism +- Worker decommissioning +""" + +import asyncio +import json +import uuid + +import pytest +import pytest_asyncio + +from hindsight_api.engine.task_backend import BrokerTaskBackend, SyncTaskBackend + + +# Use loadgroup to ensure these tests run in the same worker +# since they share database state +pytestmark = pytest.mark.xdist_group("worker_tests") + + +@pytest_asyncio.fixture +async def pool(pg0_db_url): + """Create a dedicated connection pool for worker tests.""" + import asyncpg + + from hindsight_api.pg0 import resolve_database_url + + # Resolve pg0:// URL to postgresql:// URL if needed + resolved_url = await resolve_database_url(pg0_db_url) + + pool = await asyncpg.create_pool( + resolved_url, + min_size=2, + max_size=10, + command_timeout=30, + ) + yield pool + await pool.close() + + +@pytest_asyncio.fixture +async def clean_operations(pool): + """Clean up async_operations table before and after tests.""" + # Clean before test + await pool.execute("DELETE FROM async_operations WHERE bank_id LIKE 'test-worker-%'") + yield + # Clean after test + await pool.execute("DELETE FROM async_operations WHERE bank_id LIKE 'test-worker-%'") + + +class TestBrokerTaskBackend: + """Tests for BrokerTaskBackend task storage.""" + + @pytest.mark.asyncio + async def test_submit_task_updates_existing_operation(self, pool, clean_operations): + """Test that submit_task updates task_payload for existing operations.""" + # Create an operation record first + operation_id = uuid.uuid4() + bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" + + await pool.execute( + """ + INSERT INTO async_operations (operation_id, bank_id, operation_type, status) + VALUES ($1, $2, 'test_operation', 'pending') + """, + operation_id, + bank_id, + ) + + # Submit task with same operation_id + backend = BrokerTaskBackend(pool_getter=lambda: pool) + await backend.initialize() + + task_dict = { + "operation_id": str(operation_id), + "type": "test_task", + "bank_id": bank_id, + "data": {"key": "value"}, + } + await backend.submit_task(task_dict) + + # Verify task_payload was stored + row = await pool.fetchrow( + "SELECT task_payload, status FROM async_operations WHERE operation_id = $1", + operation_id, + ) + assert row is not None + assert row["status"] == "pending" + payload = json.loads(row["task_payload"]) + assert payload["type"] == "test_task" + assert payload["data"] == {"key": "value"} + + @pytest.mark.asyncio + async def test_submit_task_creates_new_operation(self, pool, clean_operations): + """Test that submit_task creates new operation when no operation_id provided.""" + backend = BrokerTaskBackend(pool_getter=lambda: pool) + await backend.initialize() + + bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" + task_dict = { + "type": "access_count_update", + "bank_id": bank_id, + "node_ids": ["node1", "node2"], + } + await backend.submit_task(task_dict) + + # Verify new operation was created + row = await pool.fetchrow( + "SELECT operation_type, status, task_payload FROM async_operations WHERE bank_id = $1", + bank_id, + ) + assert row is not None + assert row["operation_type"] == "access_count_update" + assert row["status"] == "pending" + payload = json.loads(row["task_payload"]) + assert payload["node_ids"] == ["node1", "node2"] + + +class TestWorkerPoller: + """Tests for WorkerPoller task claiming and execution.""" + + @pytest.mark.asyncio + async def test_claim_batch_claims_pending_tasks(self, pool, clean_operations): + """Test that claim_batch claims pending tasks with task_payload.""" + from hindsight_api.worker import WorkerPoller + + # Create some pending tasks + bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" + for i in range(3): + op_id = uuid.uuid4() + payload = json.dumps({"type": "test_task", "index": i, "bank_id": bank_id}) + await pool.execute( + """ + INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload) + VALUES ($1, $2, 'test', 'pending', $3::jsonb) + """, + op_id, + bank_id, + payload, + ) + + # Create poller and claim tasks + executed_tasks = [] + + async def mock_executor(task_dict): + executed_tasks.append(task_dict) + + poller = WorkerPoller( + pool=pool, + worker_id="test-worker-1", + executor=mock_executor, + batch_size=10, + ) + + claimed = await poller.claim_batch() + assert len(claimed) == 3 + + # Verify tasks are marked as processing with worker_id + rows = await pool.fetch( + "SELECT status, worker_id FROM async_operations WHERE bank_id = $1", + bank_id, + ) + for row in rows: + assert row["status"] == "processing" + assert row["worker_id"] == "test-worker-1" + + @pytest.mark.asyncio + async def test_claim_batch_respects_batch_size(self, pool, clean_operations): + """Test that claim_batch respects the batch_size limit.""" + from hindsight_api.worker import WorkerPoller + + # Create 10 pending tasks + bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" + for i in range(10): + op_id = uuid.uuid4() + payload = json.dumps({"type": "test_task", "index": i, "bank_id": bank_id}) + await pool.execute( + """ + INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload) + VALUES ($1, $2, 'test', 'pending', $3::jsonb) + """, + op_id, + bank_id, + payload, + ) + + # Claim with batch_size=3 + poller = WorkerPoller( + pool=pool, + worker_id="test-worker-1", + executor=lambda x: None, + batch_size=3, + ) + + claimed = await poller.claim_batch() + assert len(claimed) == 3 + + @pytest.mark.asyncio + async def test_execute_task_marks_completed(self, pool, clean_operations): + """Test that successful task execution marks task as completed.""" + from hindsight_api.worker import WorkerPoller + + # Create a pending task + bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" + op_id = uuid.uuid4() + payload = json.dumps({"type": "test_task", "operation_id": str(op_id), "bank_id": bank_id}) + await pool.execute( + """ + INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload, worker_id) + VALUES ($1, $2, 'test', 'processing', $3::jsonb, 'test-worker-1') + """, + op_id, + bank_id, + payload, + ) + + executed = [] + + async def mock_executor(task_dict): + executed.append(task_dict) + + poller = WorkerPoller( + pool=pool, + worker_id="test-worker-1", + executor=mock_executor, + ) + + # Execute the task + task_dict = json.loads(payload) + await poller.execute_task(str(op_id), task_dict) + + assert len(executed) == 1 + + # Verify task is marked as completed + row = await pool.fetchrow( + "SELECT status, completed_at FROM async_operations WHERE operation_id = $1", + op_id, + ) + assert row["status"] == "completed" + assert row["completed_at"] is not None + + @pytest.mark.asyncio + async def test_execute_task_retries_on_failure(self, pool, clean_operations): + """Test that failed task execution triggers retry mechanism.""" + from hindsight_api.worker import WorkerPoller + + # Create a pending task with retry_count=0 + bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" + op_id = uuid.uuid4() + payload = json.dumps({"type": "test_task", "operation_id": str(op_id), "bank_id": bank_id}) + await pool.execute( + """ + INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload, worker_id, retry_count) + VALUES ($1, $2, 'test', 'processing', $3::jsonb, 'test-worker-1', 0) + """, + op_id, + bank_id, + payload, + ) + + async def failing_executor(task_dict): + raise ValueError("Simulated failure") + + poller = WorkerPoller( + pool=pool, + worker_id="test-worker-1", + executor=failing_executor, + max_retries=3, + ) + + # Execute (should fail and retry) + task_dict = json.loads(payload) + await poller.execute_task(str(op_id), task_dict) + + # Verify task is back to pending with incremented retry_count + row = await pool.fetchrow( + "SELECT status, retry_count, worker_id FROM async_operations WHERE operation_id = $1", + op_id, + ) + assert row["status"] == "pending" + assert row["retry_count"] == 1 + assert row["worker_id"] is None # Worker ID cleared for retry + + @pytest.mark.asyncio + async def test_execute_task_fails_after_max_retries(self, pool, clean_operations): + """Test that task is marked failed after exceeding max retries.""" + from hindsight_api.worker import WorkerPoller + + # Create a task that has already used all retries + bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" + op_id = uuid.uuid4() + payload = json.dumps({"type": "test_task", "operation_id": str(op_id), "bank_id": bank_id}) + await pool.execute( + """ + INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload, worker_id, retry_count) + VALUES ($1, $2, 'test', 'processing', $3::jsonb, 'test-worker-1', 3) + """, + op_id, + bank_id, + payload, + ) + + async def failing_executor(task_dict): + raise ValueError("Simulated failure") + + poller = WorkerPoller( + pool=pool, + worker_id="test-worker-1", + executor=failing_executor, + max_retries=3, + ) + + # Execute (should fail permanently) + task_dict = json.loads(payload) + await poller.execute_task(str(op_id), task_dict) + + # Verify task is marked as failed + row = await pool.fetchrow( + "SELECT status, error_message FROM async_operations WHERE operation_id = $1", + op_id, + ) + assert row["status"] == "failed" + assert "Max retries" in row["error_message"] + + +class TestConcurrentWorkers: + """Tests for concurrent worker task claiming (FOR UPDATE SKIP LOCKED).""" + + @pytest.mark.asyncio + async def test_concurrent_workers_claim_different_tasks(self, pool, clean_operations): + """Test that multiple workers claim different tasks (no duplicates).""" + from hindsight_api.worker import WorkerPoller + + # Create 10 pending tasks + bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" + task_ids = [] + for i in range(10): + op_id = uuid.uuid4() + task_ids.append(op_id) + payload = json.dumps({"type": "test_task", "index": i, "bank_id": bank_id, "operation_id": str(op_id)}) + await pool.execute( + """ + INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload) + VALUES ($1, $2, 'test', 'pending', $3::jsonb) + """, + op_id, + bank_id, + payload, + ) + + # Create 3 workers that will claim tasks concurrently + workers_claimed: dict[str, list[str]] = {"worker-1": [], "worker-2": [], "worker-3": []} + + async def claim_for_worker(worker_id: str): + poller = WorkerPoller( + pool=pool, + worker_id=worker_id, + executor=lambda x: None, + batch_size=5, # Each worker tries to claim 5 + ) + claimed = await poller.claim_batch() + workers_claimed[worker_id] = [op_id for op_id, _ in claimed] + + # Run all workers concurrently + await asyncio.gather( + claim_for_worker("worker-1"), + claim_for_worker("worker-2"), + claim_for_worker("worker-3"), + ) + + # Verify no duplicates - each task claimed by exactly one worker + all_claimed = workers_claimed["worker-1"] + workers_claimed["worker-2"] + workers_claimed["worker-3"] + assert len(all_claimed) == len(set(all_claimed)), "Duplicate task claimed by multiple workers!" + + # Verify total claimed equals available tasks (10) + assert len(all_claimed) == 10, f"Expected 10 tasks claimed, got {len(all_claimed)}" + + # Verify each task is assigned to exactly one worker in DB + rows = await pool.fetch( + "SELECT operation_id, worker_id FROM async_operations WHERE bank_id = $1", + bank_id, + ) + worker_assignments = {str(row["operation_id"]): row["worker_id"] for row in rows} + + # With FOR UPDATE SKIP LOCKED, it's a race condition which workers get tasks. + # The important invariant is no duplicates and all tasks claimed, which we verified above. + # Just verify that at least 1 worker got tasks and all tasks have a worker assigned. + assert len(set(worker_assignments.values())) >= 1, "At least one worker should have claimed tasks" + assert all(w is not None for w in worker_assignments.values()), "All tasks should have a worker assigned" + + @pytest.mark.asyncio + async def test_workers_do_not_claim_already_processing_tasks(self, pool, clean_operations): + """Test that workers skip tasks already being processed by another worker.""" + from hindsight_api.worker import WorkerPoller + + # Create tasks - some pending, some already processing + bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" + + # Create 3 pending tasks + for i in range(3): + op_id = uuid.uuid4() + payload = json.dumps({"type": "test_task", "index": i, "bank_id": bank_id}) + await pool.execute( + """ + INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload) + VALUES ($1, $2, 'test', 'pending', $3::jsonb) + """, + op_id, + bank_id, + payload, + ) + + # Create 2 already-processing tasks owned by another worker + for i in range(2): + op_id = uuid.uuid4() + payload = json.dumps({"type": "test_task", "index": i + 10, "bank_id": bank_id}) + await pool.execute( + """ + INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload, worker_id) + VALUES ($1, $2, 'test', 'processing', $3::jsonb, 'other-worker') + """, + op_id, + bank_id, + payload, + ) + + # New worker should only claim the 3 pending tasks + poller = WorkerPoller( + pool=pool, + worker_id="new-worker", + executor=lambda x: None, + batch_size=10, + ) + + claimed = await poller.claim_batch() + assert len(claimed) == 3, "Worker should only claim pending tasks" + + # Verify other worker's tasks are still owned by them + row = await pool.fetchrow( + "SELECT COUNT(*) as count FROM async_operations WHERE bank_id = $1 AND worker_id = 'other-worker'", + bank_id, + ) + assert row["count"] == 2 + + +class TestWorkerDecommission: + """Tests for worker decommissioning functionality.""" + + @pytest.mark.asyncio + async def test_decommission_releases_worker_tasks(self, pool, clean_operations): + """Test that decommissioning a worker releases all its processing tasks.""" + # Create tasks being processed by a worker + bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" + worker_id = "worker-to-decommission" + + for i in range(5): + op_id = uuid.uuid4() + payload = json.dumps({"type": "test_task", "index": i, "bank_id": bank_id}) + await pool.execute( + """ + INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload, worker_id, claimed_at) + VALUES ($1, $2, 'test', 'processing', $3::jsonb, $4, now()) + """, + op_id, + bank_id, + payload, + worker_id, + ) + + # Run decommission + result = await pool.fetch( + """ + UPDATE async_operations + SET status = 'pending', worker_id = NULL, claimed_at = NULL, updated_at = now() + WHERE worker_id = $1 AND status = 'processing' + RETURNING operation_id + """, + worker_id, + ) + + assert len(result) == 5 + + # Verify all tasks are back to pending + rows = await pool.fetch( + "SELECT status, worker_id, claimed_at FROM async_operations WHERE bank_id = $1", + bank_id, + ) + for row in rows: + assert row["status"] == "pending" + assert row["worker_id"] is None + assert row["claimed_at"] is None + + @pytest.mark.asyncio + async def test_decommission_does_not_affect_other_workers(self, pool, clean_operations): + """Test that decommissioning one worker doesn't affect another worker's tasks.""" + bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" + + # Create tasks for worker-1 + for i in range(3): + op_id = uuid.uuid4() + payload = json.dumps({"type": "test_task", "index": i, "bank_id": bank_id}) + await pool.execute( + """ + INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload, worker_id) + VALUES ($1, $2, 'test', 'processing', $3::jsonb, 'worker-1') + """, + op_id, + bank_id, + payload, + ) + + # Create tasks for worker-2 + for i in range(3): + op_id = uuid.uuid4() + payload = json.dumps({"type": "test_task", "index": i + 10, "bank_id": bank_id}) + await pool.execute( + """ + INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload, worker_id) + VALUES ($1, $2, 'test', 'processing', $3::jsonb, 'worker-2') + """, + op_id, + bank_id, + payload, + ) + + # Decommission worker-1 only + await pool.execute( + """ + UPDATE async_operations + SET status = 'pending', worker_id = NULL, claimed_at = NULL + WHERE worker_id = 'worker-1' AND status = 'processing' + """ + ) + + # Verify worker-1 tasks are released + worker1_rows = await pool.fetch( + "SELECT status, worker_id FROM async_operations WHERE bank_id = $1 AND worker_id IS NULL", + bank_id, + ) + assert len(worker1_rows) == 3 + + # Verify worker-2 tasks are unaffected + worker2_rows = await pool.fetch( + "SELECT status, worker_id FROM async_operations WHERE bank_id = $1 AND worker_id = 'worker-2'", + bank_id, + ) + assert len(worker2_rows) == 3 + for row in worker2_rows: + assert row["status"] == "processing" + + +class TestSyncTaskBackend: + """Tests for SyncTaskBackend (used in tests and embedded mode).""" + + @pytest.mark.asyncio + async def test_sync_backend_executes_immediately(self): + """Test that SyncTaskBackend executes tasks immediately.""" + executed = [] + + async def mock_executor(task_dict): + executed.append(task_dict) + + backend = SyncTaskBackend() + backend.set_executor(mock_executor) + await backend.initialize() + + task_dict = {"type": "test", "data": "value"} + await backend.submit_task(task_dict) + + assert len(executed) == 1 + assert executed[0] == task_dict + + @pytest.mark.asyncio + async def test_sync_backend_handles_errors(self): + """Test that SyncTaskBackend handles executor errors gracefully.""" + + async def failing_executor(task_dict): + raise ValueError("Test error") + + backend = SyncTaskBackend() + backend.set_executor(failing_executor) + await backend.initialize() + + # Should not raise, error is logged + await backend.submit_task({"type": "test"}) diff --git a/hindsight-docs/docs/developer/admin-cli.md b/hindsight-docs/docs/developer/admin-cli.md index eee17675..98491b7c 100644 --- a/hindsight-docs/docs/developer/admin-cli.md +++ b/hindsight-docs/docs/developer/admin-cli.md @@ -127,6 +127,53 @@ Restore will **delete all existing data** in the target schema before importing --- +### decommission-worker + +Release all tasks owned by a worker, resetting them from "processing" back to "pending" status so they can be picked up by other workers. + +```bash +hindsight-admin decommission-worker WORKER_ID [OPTIONS] +``` + +**Arguments:** + +| Argument | Description | +|----------|-------------| +| `WORKER_ID` | ID of the worker to decommission | + +**Options:** + +| Option | Description | Default | +|--------|-------------|---------| +| `--schema`, `-s` | Database schema | `public` | + +**Examples:** + +```bash +# Before scaling down - release tasks from workers being removed +hindsight-admin decommission-worker hindsight-worker-4 +hindsight-admin decommission-worker hindsight-worker-3 + +# Release tasks from a crashed worker +hindsight-admin decommission-worker worker-2 + +# For a specific tenant schema +hindsight-admin decommission-worker worker-1 --schema tenant_acme +``` + +**When to Use:** + +- **Scaling down**: Before removing worker replicas in Kubernetes +- **Graceful removal**: When taking a worker offline for maintenance +- **Crash recovery**: If a worker crashed while processing tasks +- **Stuck worker**: When a worker is unresponsive + +:::tip Finding Worker IDs +Worker IDs default to the hostname. In Kubernetes StatefulSets, this is the pod name (e.g., `hindsight-worker-0`). You can also set a custom ID with `HINDSIGHT_API_WORKER_ID` or `--worker-id`. +::: + +--- + ## Environment Variables The admin CLI uses the same environment variables as the API service. The most important one is: diff --git a/hindsight-docs/docs/developer/configuration.md b/hindsight-docs/docs/developer/configuration.md index 26d6e215..23ac70af 100644 --- a/hindsight-docs/docs/developer/configuration.md +++ b/hindsight-docs/docs/developer/configuration.md @@ -344,15 +344,18 @@ Configuration for the local MCP server (`hindsight-local-mcp` command). export HINDSIGHT_API_MCP_INSTRUCTIONS="Also store every action you take, including tool calls and decisions made." ``` -### Background Tasks +### Distributed Workers -Controls background task processing for async operations like opinion formation and entity observations. +Configuration for background task processing. By default, the API processes tasks internally. For high-throughput deployments, run dedicated workers. See [Services - Worker Service](./services#worker-service) for details. | Variable | Description | Default | |----------|-------------|---------| -| `HINDSIGHT_API_TASK_BACKEND` | Task backend implementation: `memory` (in-process queue) or `noop` (discard tasks, useful for tests) | `memory` | -| `HINDSIGHT_API_TASK_BACKEND_MEMORY_BATCH_SIZE` | Max tasks to process in one batch (memory backend only) | `10` | -| `HINDSIGHT_API_TASK_BACKEND_MEMORY_BATCH_INTERVAL` | Interval between batch processing in seconds (memory backend only) | `1.0` | +| `HINDSIGHT_API_WORKER_ENABLED` | Enable internal worker in API process | `true` | +| `HINDSIGHT_API_WORKER_ID` | Unique worker identifier | hostname | +| `HINDSIGHT_API_WORKER_POLL_INTERVAL_MS` | Database polling interval in milliseconds | `500` | +| `HINDSIGHT_API_WORKER_BATCH_SIZE` | Tasks to claim per poll cycle | `10` | +| `HINDSIGHT_API_WORKER_MAX_RETRIES` | Max retries before marking task failed | `3` | +| `HINDSIGHT_API_WORKER_HTTP_PORT` | HTTP port for worker metrics/health (worker CLI only) | `8889` | ### Performance Optimization diff --git a/hindsight-docs/docs/developer/installation.md b/hindsight-docs/docs/developer/installation.md index 31b28ae3..eea36344 100644 --- a/hindsight-docs/docs/developer/installation.md +++ b/hindsight-docs/docs/developer/installation.md @@ -36,8 +36,6 @@ See [Models](./models) for detailed comparison and configuration. **Best for**: Quick start, development, small deployments -### Single Container (Quickest) - Run everything in one container with embedded PostgreSQL: ```bash @@ -83,7 +81,19 @@ helm upgrade hindsight oci://ghcr.io/vectorize-io/charts/hindsight - Kubernetes cluster (GKE, EKS, AKS, or self-hosted) - Helm 3.8+ -See the [Helm chart documentation](https://github.com/vectorize-io/hindsight/tree/main/helm) for advanced configuration. +### Distributed Workers + +For high-throughput deployments, enable dedicated worker pods to scale task processing independently: + +```bash +helm install hindsight oci://ghcr.io/vectorize-io/charts/hindsight \ + --set worker.enabled=true \ + --set worker.replicaCount=3 +``` + +See [Services - Worker Service](./services#worker-service) for configuration details and architecture. + +See the [Helm chart values.yaml](https://github.com/vectorize-io/hindsight/tree/main/helm/hindsight/values.yaml) for all chart options. --- diff --git a/hindsight-docs/docs/developer/services.md b/hindsight-docs/docs/developer/services.md index 4b89c2de..2dc19976 100644 --- a/hindsight-docs/docs/developer/services.md +++ b/hindsight-docs/docs/developer/services.md @@ -1,6 +1,6 @@ # Services -Hindsight consists of two services that can run together or separately depending on your deployment needs. +Hindsight consists of three services that can run together or separately depending on your deployment needs. ## API Service @@ -10,12 +10,48 @@ The core memory engine. Handles all memory operations: - **Recall**: Semantic search across memories - **Reflect**: Disposition-aware answer generation -``` +```bash hindsight-api # Default port: 8888 ``` The API service is stateless and can be horizontally scaled behind a load balancer. All state is stored in PostgreSQL. +By default, the API also processes background tasks (opinion formation, entity observations) internally. For high-throughput deployments, you can disable this and run dedicated workers instead. + +## Worker Service + +Dedicated task processor for background operations. Uses the **same package and Docker image** as the API service, just with a different entry point. + +```bash +hindsight-worker # Default metrics port: 8889 +``` + +Workers use PostgreSQL as a task broker, polling for pending tasks. Multiple workers can run simultaneously without conflicts. + +| Deployment | Internal Worker | Dedicated Workers | +|------------|-----------------|-------------------| +| **Development** | ✅ Simple, all-in-one | ❌ Overkill | +| **Small production** | ✅ Less infrastructure | ❌ Overkill | +| **High throughput** | ❌ API bottleneck | ✅ Scale independently | +| **Long-running tasks** | ❌ Blocks API resources | ✅ Isolated processing | + +To use dedicated workers, disable the internal worker in the API and start worker processes: + +```bash +# Disable internal worker in API +HINDSIGHT_API_WORKER_ENABLED=false hindsight-api + +# Start dedicated workers (run multiple instances) +hindsight-worker --worker-id worker-1 +hindsight-worker --worker-id worker-2 +``` + +Each worker exposes `/health` and `/metrics` endpoints for monitoring. + +Before scaling down or removing workers, release their tasks with `hindsight-admin decommission-worker `. + +See [Configuration - Distributed Workers](./configuration#distributed-workers) for all worker settings and [Installation - Helm](./installation#distributed-workers) for Kubernetes deployment. + ## Control Plane Web UI for managing and exploring your memory banks: @@ -28,13 +64,3 @@ Web UI for managing and exploring your memory banks: The Control Plane connects to the API service and provides a visual interface for development and debugging. For bare metal deployments, you can run the Control Plane standalone using npx. See [Installation - Bare Metal](./installation#control-plane) for details. - -## Deployment Options - -| Deployment | Services | Use Case | -|------------|----------|----------| -| **Docker (single container)** | Both bundled | Development, quick start | -| **Helm / Kubernetes** | Separate pods | Production, scaling | -| **Bare metal** | Run independently | Custom deployments | - -In the Docker quickstart, both services run in a single container. For production Kubernetes deployments, they run as separate pods with independent scaling. For bare metal, you can run the API via pip and the Control Plane via npx. diff --git a/scripts/dev/start-worker.sh b/scripts/dev/start-worker.sh new file mode 100755 index 00000000..4df589b8 --- /dev/null +++ b/scripts/dev/start-worker.sh @@ -0,0 +1,20 @@ +#!/bin/bash +set -e + +cd "$(dirname "$0")/../.." + +ENV_FILE=".env" +if [ ! -f "$ENV_FILE" ]; then + echo "Error: Environment file $ENV_FILE not found at project root." + exit 1 +fi + +echo "Loading environment from $ENV_FILE" +echo "" + +# Export all variables from env file +set -a +source "$ENV_FILE" +set +a + +uv run hindsight-worker "$@"