Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
85d5b3bfd9 | ||
|
|
f15b76fba8 | ||
|
|
61f457736e | ||
|
|
5de1447ff5 |
@@ -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
|
||||
*/}}
|
||||
|
||||
@@ -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 }}
|
||||
|
||||
@@ -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 }}
|
||||
@@ -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 }}
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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()
|
||||
|
||||
@@ -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))),
|
||||
)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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"]
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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"
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"})
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -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 <worker-id>`.
|
||||
|
||||
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.
|
||||
|
||||
Executable
+20
@@ -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 "$@"
|
||||
Reference in New Issue
Block a user