e588d97e0f
- Updated .env.example and config.yaml to include observability settings. - Added a new Phoenix service in docker-compose.yml for self-hosted tracing. - Enhanced README.md with instructions on enabling and using observability features. - Implemented tracing in the pipeline, including job spans and LLM stage spans. - Introduced ObservabilitySettings class in config.py for better configuration management. - Updated job.py and resources.py to support tracing without affecting fault isolation. - Minor adjustments to other files for compatibility with the new observability features.
228 lines
9.2 KiB
Python
228 lines
9.2 KiB
Python
"""Resources shared across every job in a batch run.
|
|
|
|
A single-field CLI run opened one SQL connection, one Weaviate client, and
|
|
made its two LLM calls without needing to coordinate with anyone else. The
|
|
batch runs many jobs concurrently across a thread pool, so several things
|
|
that used to be "create it, use it, close it" per field now need to be
|
|
long-lived, shared, and safe to touch from multiple threads at once:
|
|
|
|
- pyodbc connections must not be shared across threads, so each worker thread
|
|
gets its own, created lazily and reused for every job it picks up.
|
|
- The Weaviate client is documented as safe for concurrent query use, so one
|
|
instance is opened for the whole batch instead of one per field.
|
|
- `near_text` search embeds through the same Gemini API key as the two LLM
|
|
calls (`ProductProfile` is vectorised with text2vec-palm /
|
|
gemini-embedding-001), so a single semaphore + rate limiter covers all three
|
|
Gemini-backed calls per job instead of treating "LLM" and "Weaviate" as
|
|
independent budgets.
|
|
- The `products` table and the `ProductJsonBuilder` per-product cache are
|
|
loaded/shared once instead of once per field.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import threading
|
|
import time
|
|
from typing import Any
|
|
|
|
from pipeline.config import LlmSettings, Settings
|
|
from pipeline.db import connect_raw
|
|
from pipeline.observability import llm_stage_span, retrieval_span, set_retrieval_results
|
|
from pipeline.prompts import PromptRegistry
|
|
from pipeline.retry import call_with_llm_retry, call_with_weaviate_retry
|
|
from pipeline.stages.llm import call_llm as _raw_call_llm
|
|
from pipeline.stages.product_json import ProductJsonBuilder
|
|
from pipeline.stages.products import ProductCandidate, ProductIndex
|
|
from pipeline.stages.vector import RecommendedProduct, connect_client
|
|
from pipeline.stages.vector import search_products as _raw_search_products
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class RateLimiter:
|
|
"""
|
|
Thread-safe pacing limiter: spreads calls out to at most N per minute.
|
|
|
|
This paces call *starts* evenly (a leaky bucket, not a bursty token
|
|
bucket with saved-up credit) — the simplest thing that reliably keeps a
|
|
pool of worker threads under an external rate limit, which matters more
|
|
here than allowing occasional bursts.
|
|
"""
|
|
|
|
def __init__(self, requests_per_minute: int) -> None:
|
|
self._interval = 60.0 / requests_per_minute if requests_per_minute > 0 else 0.0
|
|
self._lock = threading.Lock()
|
|
self._next_allowed = 0.0
|
|
|
|
def acquire(self) -> None:
|
|
if self._interval <= 0:
|
|
return
|
|
with self._lock:
|
|
now = time.monotonic()
|
|
wait = self._next_allowed - now
|
|
if wait > 0:
|
|
time.sleep(wait)
|
|
now = time.monotonic()
|
|
self._next_allowed = max(now, self._next_allowed) + self._interval
|
|
|
|
|
|
def _llm_stage_name(settings: Settings, llm_cfg: LlmSettings) -> str:
|
|
"""Identify which of the two configured LLM profiles a call belongs to,
|
|
for the tracing span name (see `Resources.call_llm`)."""
|
|
if llm_cfg is settings.llm_query_synthesis:
|
|
return "query_synthesis"
|
|
if llm_cfg is settings.llm_advice_generation:
|
|
return "advice_generation"
|
|
return "llm"
|
|
|
|
|
|
def _llm_model_name(llm_cfg: LlmSettings) -> str:
|
|
return {
|
|
"gemini": llm_cfg.gemini_model,
|
|
"openai": llm_cfg.openai_model,
|
|
"anthropic": llm_cfg.anthropic_model,
|
|
}.get(llm_cfg.provider, "?")
|
|
|
|
|
|
class Resources:
|
|
"""Shared, thread-safe state for one batch run. Create once, close once."""
|
|
|
|
def __init__(self, settings: Settings, prompts: PromptRegistry) -> None:
|
|
self.settings = settings
|
|
self.prompts = prompts
|
|
|
|
self._local = threading.local()
|
|
self._connections: list[Any] = []
|
|
self._connections_lock = threading.Lock()
|
|
self.sql_sem = threading.Semaphore(settings.concurrency.sql_limit)
|
|
|
|
self.gemini_sem = threading.Semaphore(settings.concurrency.gemini_limit)
|
|
self.gemini_rate_limiter = RateLimiter(settings.concurrency.gemini_requests_per_minute)
|
|
|
|
self._weaviate_client: Any = None
|
|
self._weaviate_lock = threading.Lock()
|
|
|
|
self.product_index: ProductIndex | None = None
|
|
|
|
self._builder_cache: dict[tuple[str, str], ProductJsonBuilder] = {}
|
|
self._builder_lock = threading.Lock()
|
|
|
|
# -- SQL -----------------------------------------------------------------
|
|
|
|
def sql_connection(self) -> Any:
|
|
"""
|
|
Return this thread's SQL connection, opening one on first use.
|
|
|
|
Not gated by `sql_sem`: with one connection per worker thread kept
|
|
open for the batch's whole lifetime, gating *creation* would only
|
|
throttle how fast the pool warms up, and gating it for the
|
|
connection's entire lifetime would deadlock as soon as `sql_limit` is
|
|
set below `workers` (the remaining threads would block forever on a
|
|
permit nothing ever releases). `sql_sem` instead bounds how many jobs
|
|
may run their SQL-heavy phase at the same instant — see
|
|
`pipeline.job.run_job`, which acquires it around each phase, not
|
|
around the connection.
|
|
"""
|
|
conn = getattr(self._local, "conn", None)
|
|
if conn is None:
|
|
conn = connect_raw(self.settings.sql)
|
|
self._local.conn = conn
|
|
with self._connections_lock:
|
|
self._connections.append(conn)
|
|
return conn
|
|
|
|
def load_product_index(self) -> ProductIndex:
|
|
"""Load the whole `products` table once; call before starting the pool."""
|
|
self.product_index = ProductIndex.load(self.sql_connection())
|
|
return self.product_index
|
|
|
|
def product_json_builder(self, crop_english: str, disease_english: str) -> ProductJsonBuilder:
|
|
"""
|
|
One `ProductJsonBuilder` per (crop, disease), shared by every field
|
|
growing that crop with that disease across all worker threads.
|
|
"""
|
|
if self.product_index is None:
|
|
raise RuntimeError("load_product_index() must run before product_json_builder().")
|
|
key = (crop_english, disease_english)
|
|
with self._builder_lock:
|
|
builder = self._builder_cache.get(key)
|
|
if builder is None:
|
|
builder = ProductJsonBuilder(self.product_index, crop_english, disease_english)
|
|
self._builder_cache[key] = builder
|
|
return builder
|
|
|
|
# -- Weaviate / Gemini -----------------------------------------------------
|
|
|
|
def _weaviate(self) -> Any:
|
|
if self._weaviate_client is None:
|
|
with self._weaviate_lock:
|
|
if self._weaviate_client is None:
|
|
self._weaviate_client = connect_client(self.settings)
|
|
return self._weaviate_client
|
|
|
|
def call_llm(
|
|
self,
|
|
settings: Settings,
|
|
llm_cfg: LlmSettings,
|
|
system_prompt: str,
|
|
user_content: str,
|
|
json_output: bool = False,
|
|
) -> str:
|
|
"""
|
|
Rate-limited, retried, concurrency-bounded drop-in for
|
|
`pipeline.stages.llm.call_llm`.
|
|
|
|
Signature-compatible so it can be passed as `synthesize_queries(...,
|
|
llm_call=resources.call_llm)` / `generate_advice(...,
|
|
llm_call=resources.call_llm)`. Wrapped in a tracing span (a no-op
|
|
when observability is disabled) named after the calling stage, so a
|
|
retried call still produces one stage span with the Gemini/OpenAI/
|
|
Anthropic auto-instrumentor's own LLM span nested inside it.
|
|
"""
|
|
|
|
def _attempt() -> str:
|
|
with self.gemini_sem:
|
|
self.gemini_rate_limiter.acquire()
|
|
return _raw_call_llm(settings, llm_cfg, system_prompt, user_content, json_output)
|
|
|
|
stage = _llm_stage_name(self.settings, llm_cfg)
|
|
with llm_stage_span(stage, llm_cfg.provider, _llm_model_name(llm_cfg)):
|
|
return call_with_llm_retry(self.settings.retry, _attempt)
|
|
|
|
def search_products(
|
|
self,
|
|
queries: list[str],
|
|
candidates: list[ProductCandidate],
|
|
) -> list[RecommendedProduct]:
|
|
"""Rate-limited, retried wrapper around `pipeline.stages.vector.search_products`."""
|
|
|
|
def _attempt() -> list[RecommendedProduct]:
|
|
with self.gemini_sem:
|
|
self.gemini_rate_limiter.acquire()
|
|
return _raw_search_products(self._weaviate(), queries, candidates)
|
|
|
|
with retrieval_span(queries, len(candidates)) as span:
|
|
results = call_with_weaviate_retry(self.settings.retry, _attempt)
|
|
set_retrieval_results(span, results)
|
|
return results
|
|
|
|
# -- Lifecycle -------------------------------------------------------------
|
|
|
|
def close(self) -> None:
|
|
"""Close every pooled SQL connection and the shared Weaviate client."""
|
|
with self._connections_lock:
|
|
for conn in self._connections:
|
|
try:
|
|
conn.close()
|
|
except Exception:
|
|
logger.warning("Error closing a pooled SQL connection.", exc_info=True)
|
|
self._connections.clear()
|
|
|
|
if self._weaviate_client is not None:
|
|
try:
|
|
self._weaviate_client.close()
|
|
except Exception:
|
|
logger.warning("Error closing the Weaviate client.", exc_info=True)
|
|
self._weaviate_client = None
|