Rewrite job_runner, event_bridge, cost estimator, and DigiKey client.
Local pid-file dispatch stays the self-host path. DigiKey still rejects unrelated MPNs. Inherited copies remain on disk.
This commit is contained in:
@@ -0,0 +1,272 @@
|
||||
"""Read-only pre-flight USD estimate from BOM + library cache. DeepSeek rates."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import Literal
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from backend.config import settings
|
||||
from backend.periscopex.parsers import parse_bom
|
||||
from backend.periscopex.resolve_passives import resolve_mpn
|
||||
from backend.periscopex.taxonomy import SIMPLE_TYPES, type_for_ref
|
||||
from backend.periscopex.utils import safe_mpn
|
||||
from backend.services import projects as proj_svc
|
||||
from backend.services.billing_hook import get_billing
|
||||
from backend.services.llm.pricing import CACHE_RATES, PRICING
|
||||
from backend.services.storage import LocalStorageBackend, StorageBackend
|
||||
|
||||
STAGE_TOKEN_BASELINES: dict[str, dict[str, int | str]] = {
|
||||
"ic_extraction": {
|
||||
"settings_stage": "pintable",
|
||||
"input": 100, "output": 2000,
|
||||
"cache_create": 80_000, "cache_read": 170_000,
|
||||
},
|
||||
"simple_extraction": {
|
||||
"settings_stage": "specs",
|
||||
"input": 100, "output": 1000,
|
||||
"cache_create": 20_000, "cache_read": 20_000,
|
||||
},
|
||||
"passive_pattern": {
|
||||
"settings_stage": "pattern",
|
||||
"input": 100, "output": 7000,
|
||||
"cache_create": 60_000, "cache_read": 330_000,
|
||||
},
|
||||
"digikey_resolve": {
|
||||
"settings_stage": "auto_resolve",
|
||||
"input": 2000, "output": 200,
|
||||
"cache_create": 0, "cache_read": 0,
|
||||
},
|
||||
"review": {
|
||||
"settings_stage": "validation",
|
||||
"input": 13_500, "output": 2000,
|
||||
"cache_create": 110_000, "cache_read": 300_000,
|
||||
},
|
||||
"normalize": {
|
||||
"settings_stage": "normalize",
|
||||
"input": 1500, "output": 600,
|
||||
"cache_create": 0, "cache_read": 0,
|
||||
},
|
||||
"cross_ic_dedup": {
|
||||
"settings_stage": "normalize",
|
||||
"input": 2500, "output": 700,
|
||||
"cache_create": 0, "cache_read": 0,
|
||||
},
|
||||
}
|
||||
|
||||
LOW_MULT = 0.7
|
||||
HIGH_MULT = 1.4
|
||||
|
||||
UnitKind = Literal[
|
||||
"ic_extraction",
|
||||
"simple_extraction",
|
||||
"passive_pattern",
|
||||
"digikey_resolve",
|
||||
"review",
|
||||
]
|
||||
|
||||
|
||||
class CostItem(BaseModel):
|
||||
identifier: str
|
||||
kind: UnitKind
|
||||
api_cost_usd: float
|
||||
source: Literal["cache_hit", "api_call", "api_call_estimated"]
|
||||
note: str | None = None
|
||||
|
||||
|
||||
class CostEstimate(BaseModel):
|
||||
api_cost_low: float
|
||||
api_cost_high: float
|
||||
api_cost_mid: float
|
||||
credits_low: float
|
||||
credits_high: float
|
||||
credits_mid: float
|
||||
breakdown: list[CostItem]
|
||||
ic_count: int
|
||||
simple_count: int
|
||||
passive_count: int
|
||||
cached_ic_count: int
|
||||
cached_simple_count: int
|
||||
cached_passive_count: int
|
||||
review_ic_count: int
|
||||
|
||||
|
||||
def estimate_stage_cost_usd(stage: str) -> float:
|
||||
base = STAGE_TOKEN_BASELINES[stage]
|
||||
settings_stage = str(base["settings_stage"])
|
||||
provider = settings.provider_for_stage(settings_stage)
|
||||
model = settings.model_for_stage(settings_stage)
|
||||
table = PRICING.get(provider) or PRICING["deepseek"]
|
||||
rates = table.get(model, table["default"])
|
||||
cache = CACHE_RATES.get(provider, CACHE_RATES["deepseek"])
|
||||
return (
|
||||
int(base["input"]) * rates["input"]
|
||||
+ int(base["output"]) * rates["output"]
|
||||
+ int(base["cache_create"]) * rates["input"] * cache["create"]
|
||||
+ int(base["cache_read"]) * rates["input"] * cache["read"]
|
||||
) / 1_000_000
|
||||
|
||||
|
||||
def _load_library_patterns(storage: StorageBackend):
|
||||
try:
|
||||
return proj_svc.load_library_patterns(storage)
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
|
||||
def _locate_datasheet_local(
|
||||
storage: StorageBackend, user_id: str, project_id: str, mpn: str,
|
||||
) -> Path | None:
|
||||
if not isinstance(storage, LocalStorageBackend):
|
||||
return None
|
||||
safe = safe_mpn(mpn)
|
||||
key = f"users/{user_id}/projects/{project_id}/uploads/datasheets/{safe}.pdf"
|
||||
if storage.exists(key):
|
||||
return storage._path(key) # type: ignore[attr-defined]
|
||||
legacy = f"library/datasheets/{safe}.pdf"
|
||||
if storage.exists(legacy):
|
||||
return storage._path(legacy) # type: ignore[attr-defined]
|
||||
return None
|
||||
|
||||
|
||||
def estimate_pipeline_cost(
|
||||
storage: StorageBackend, user_id: str, project_id: str,
|
||||
) -> CostEstimate:
|
||||
bom_key = proj_svc.get_bom_key(storage, user_id, project_id)
|
||||
if not bom_key:
|
||||
raise FileNotFoundError("BOM not uploaded for this project")
|
||||
meta = proj_svc.get_project(storage, user_id, project_id)
|
||||
col_map = (meta.bom_columns if meta else None) or {}
|
||||
ref_col = col_map.get("reference", "Reference")
|
||||
mpn_col = col_map.get("mpn", "Manufacturer Part Number")
|
||||
with tempfile.NamedTemporaryFile(suffix=".csv", delete=False) as tmp:
|
||||
tmp.write(storage.read_bytes(bom_key))
|
||||
bom_local_path = Path(tmp.name)
|
||||
try:
|
||||
bom = parse_bom(str(bom_local_path), reference_col=ref_col, mpn_col=mpn_col)
|
||||
finally:
|
||||
bom_local_path.unlink(missing_ok=True)
|
||||
|
||||
ic_mpns: set[str] = set()
|
||||
simple_mpns: set[str] = set()
|
||||
passive_mpns: set[str] = set()
|
||||
for ref, info in bom.items():
|
||||
mpn = info.get("mpn")
|
||||
if not mpn:
|
||||
continue
|
||||
typ = type_for_ref(ref)
|
||||
if typ == "ic":
|
||||
ic_mpns.add(mpn)
|
||||
elif typ == "passive":
|
||||
passive_mpns.add(mpn)
|
||||
elif typ in SIMPLE_TYPES:
|
||||
simple_mpns.add(mpn)
|
||||
|
||||
patterns = _load_library_patterns(storage)
|
||||
breakdown: list[CostItem] = []
|
||||
cached_ic = cached_simple = cached_passive = 0
|
||||
|
||||
for mpn in sorted(ic_mpns):
|
||||
if proj_svc.library_has_extraction(storage, mpn):
|
||||
breakdown.append(CostItem(
|
||||
identifier=mpn, kind="ic_extraction",
|
||||
api_cost_usd=0.0, source="cache_hit", note="library hit",
|
||||
))
|
||||
cached_ic += 1
|
||||
else:
|
||||
breakdown.append(CostItem(
|
||||
identifier=mpn, kind="ic_extraction",
|
||||
api_cost_usd=round(estimate_stage_cost_usd("ic_extraction"), 4),
|
||||
source="api_call_estimated",
|
||||
))
|
||||
|
||||
for mpn in sorted(simple_mpns):
|
||||
if proj_svc.library_has_model(storage, mpn):
|
||||
breakdown.append(CostItem(
|
||||
identifier=mpn, kind="simple_extraction",
|
||||
api_cost_usd=0.0, source="cache_hit",
|
||||
))
|
||||
cached_simple += 1
|
||||
else:
|
||||
breakdown.append(CostItem(
|
||||
identifier=mpn, kind="simple_extraction",
|
||||
api_cost_usd=round(estimate_stage_cost_usd("simple_extraction"), 4),
|
||||
source="api_call_estimated",
|
||||
))
|
||||
|
||||
unresolved_passives: list[str] = []
|
||||
for mpn in sorted(passive_mpns):
|
||||
if patterns and resolve_mpn(mpn, patterns) is not None:
|
||||
breakdown.append(CostItem(
|
||||
identifier=mpn, kind="passive_pattern",
|
||||
api_cost_usd=0.0, source="cache_hit", note="pattern match",
|
||||
))
|
||||
cached_passive += 1
|
||||
continue
|
||||
if proj_svc.library_has_passive_model(storage, mpn):
|
||||
breakdown.append(CostItem(
|
||||
identifier=mpn, kind="passive_pattern",
|
||||
api_cost_usd=0.0, source="cache_hit", note="cached passive model",
|
||||
))
|
||||
cached_passive += 1
|
||||
continue
|
||||
unresolved_passives.append(mpn)
|
||||
|
||||
prefixes = {m[:7] for m in unresolved_passives}
|
||||
for prefix in sorted(prefixes):
|
||||
sample_mpn = next(m for m in unresolved_passives if m.startswith(prefix))
|
||||
breakdown.append(CostItem(
|
||||
identifier=sample_mpn, kind="passive_pattern",
|
||||
api_cost_usd=round(estimate_stage_cost_usd("passive_pattern"), 4),
|
||||
source="api_call_estimated",
|
||||
note=f"may cover {sum(1 for m in unresolved_passives if m.startswith(prefix))} MPNs",
|
||||
))
|
||||
|
||||
review_per_ic = estimate_stage_cost_usd("review")
|
||||
if settings.normalize_findings_enabled:
|
||||
review_per_ic += estimate_stage_cost_usd("normalize")
|
||||
review_ic_count = 0
|
||||
for mpn in sorted(ic_mpns):
|
||||
has_pdf = _locate_datasheet_local(storage, user_id, project_id, mpn) is not None
|
||||
has_library_pdf = (
|
||||
proj_svc.library_has_datasheet(storage, mpn) is not None if not has_pdf else False
|
||||
)
|
||||
if not (has_pdf or has_library_pdf):
|
||||
continue
|
||||
breakdown.append(CostItem(
|
||||
identifier=mpn, kind="review",
|
||||
api_cost_usd=round(review_per_ic, 4),
|
||||
source="api_call_estimated",
|
||||
))
|
||||
review_ic_count += 1
|
||||
|
||||
if settings.cross_ic_dedup_enabled and review_ic_count > 1:
|
||||
breakdown.append(CostItem(
|
||||
identifier="cross-IC dedup", kind="review",
|
||||
api_cost_usd=round(estimate_stage_cost_usd("cross_ic_dedup"), 4),
|
||||
source="api_call_estimated",
|
||||
note="collapses one interface defect reported from both ICs",
|
||||
))
|
||||
|
||||
api_total = sum(item.api_cost_usd for item in breakdown)
|
||||
api_low = round(api_total * LOW_MULT, 4)
|
||||
api_high = round(api_total * HIGH_MULT, 4)
|
||||
billing = get_billing()
|
||||
return CostEstimate(
|
||||
api_cost_low=api_low,
|
||||
api_cost_high=api_high,
|
||||
api_cost_mid=round(api_total, 4),
|
||||
credits_low=billing.credits_for_api_cost(api_low),
|
||||
credits_high=billing.credits_for_api_cost(api_high),
|
||||
credits_mid=billing.credits_for_api_cost(api_total),
|
||||
breakdown=breakdown,
|
||||
ic_count=len(ic_mpns),
|
||||
simple_count=len(simple_mpns),
|
||||
passive_count=len(passive_mpns),
|
||||
cached_ic_count=cached_ic,
|
||||
cached_simple_count=cached_simple,
|
||||
cached_passive_count=cached_passive,
|
||||
review_ic_count=review_ic_count,
|
||||
)
|
||||
@@ -0,0 +1,259 @@
|
||||
"""DigiKey Product Information v4. Exact-MPN match only — never products[0]."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
import httpx
|
||||
|
||||
from backend.config import settings
|
||||
from backend.services.datasheet_finder import (
|
||||
_alnum,
|
||||
mpn_catalog_match,
|
||||
mpn_matches,
|
||||
mpn_query_variants,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_token_cache: dict[str, str | float] = {"access_token": "", "expires_at": 0.0}
|
||||
_BASE_URLS = {
|
||||
"production": "https://api.digikey.com",
|
||||
"sandbox": "https://sandbox-api.digikey.com",
|
||||
}
|
||||
_PDF_MAGIC = b"%PDF-"
|
||||
_MIN_PDF_SIZE = 5_000
|
||||
|
||||
|
||||
async def _get_access_token() -> str:
|
||||
now = time.time()
|
||||
if _token_cache["access_token"] and float(_token_cache["expires_at"]) > now + 60:
|
||||
return str(_token_cache["access_token"])
|
||||
base = _BASE_URLS.get(settings.digikey_environment, _BASE_URLS["production"])
|
||||
async with httpx.AsyncClient(timeout=15) as client:
|
||||
resp = await client.post(
|
||||
f"{base}/v1/oauth2/token",
|
||||
data={
|
||||
"client_id": settings.digikey_client_id,
|
||||
"client_secret": settings.digikey_client_secret,
|
||||
"grant_type": "client_credentials",
|
||||
},
|
||||
)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
_token_cache["access_token"] = data["access_token"]
|
||||
_token_cache["expires_at"] = now + data.get("expires_in", 3600)
|
||||
return str(_token_cache["access_token"])
|
||||
|
||||
|
||||
def _get_mpn(product: dict) -> str:
|
||||
return product.get("ManufacturerProductNumber") or product.get("ManufacturerPartNumber") or ""
|
||||
|
||||
|
||||
def _get_ds_url(product: dict) -> str:
|
||||
url = product.get("DatasheetUrl") or product.get("PrimaryDatasheet") or ""
|
||||
if url.startswith("//"):
|
||||
url = "https:" + url
|
||||
return url
|
||||
|
||||
|
||||
async def _keyword_search(mpn: str) -> list[dict]:
|
||||
base = _BASE_URLS.get(settings.digikey_environment, _BASE_URLS["production"])
|
||||
token = await _get_access_token()
|
||||
headers = {
|
||||
"Authorization": f"Bearer {token}",
|
||||
"X-DIGIKEY-Client-Id": settings.digikey_client_id,
|
||||
"X-DIGIKEY-Locale-Site": settings.digikey_locale_site,
|
||||
"X-DIGIKEY-Locale-Language": settings.digikey_locale_language,
|
||||
"X-DIGIKEY-Locale-Currency": settings.digikey_locale_currency,
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
async with httpx.AsyncClient(timeout=20) as client:
|
||||
resp = await client.post(
|
||||
f"{base}/products/v4/search/keyword",
|
||||
headers=headers,
|
||||
json={"Keywords": mpn, "Limit": 5, "Offset": 0, "ExcludeMarketPlaceProducts": True},
|
||||
)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
return data.get("Products") or data.get("products") or []
|
||||
|
||||
|
||||
def _find_product_one(mpn: str, products: list[dict]) -> dict | None:
|
||||
exact = loose = family = None
|
||||
want = _alnum(mpn)
|
||||
for product in products:
|
||||
cand = _get_mpn(product)
|
||||
if not cand:
|
||||
continue
|
||||
got = _alnum(cand)
|
||||
if got == want:
|
||||
exact = product
|
||||
break
|
||||
if loose is None and mpn_matches(mpn, cand):
|
||||
loose = product
|
||||
elif family is None and mpn_catalog_match(mpn, cand):
|
||||
family = product
|
||||
return exact or loose or family
|
||||
|
||||
|
||||
def _find_product(mpn: str, products: list[dict]) -> dict | None:
|
||||
if not products:
|
||||
return None
|
||||
for query in mpn_query_variants(mpn):
|
||||
hit = _find_product_one(query, products)
|
||||
if hit:
|
||||
return hit
|
||||
return None
|
||||
|
||||
|
||||
async def _search_mpn(mpn: str) -> tuple[str | None, str | None]:
|
||||
tried: set[str] = set()
|
||||
for keyword in mpn_query_variants(mpn):
|
||||
key = keyword.upper()
|
||||
if key in tried:
|
||||
continue
|
||||
tried.add(key)
|
||||
products = await _keyword_search(keyword)
|
||||
product = _find_product(mpn, products)
|
||||
if not product:
|
||||
continue
|
||||
url = _get_ds_url(product)
|
||||
if url:
|
||||
return url, _get_mpn(product) or None
|
||||
return None, None
|
||||
|
||||
|
||||
async def _download_pdf(url: str) -> bytes:
|
||||
async with httpx.AsyncClient(timeout=15, follow_redirects=True) as client:
|
||||
resp = await client.get(url)
|
||||
resp.raise_for_status()
|
||||
data = resp.content
|
||||
if not data.startswith(_PDF_MAGIC):
|
||||
raise ValueError("Downloaded file is not a valid PDF (bad magic bytes)")
|
||||
if len(data) < _MIN_PDF_SIZE:
|
||||
raise ValueError(f"PDF too small ({len(data)} bytes) — likely an error page")
|
||||
return data
|
||||
|
||||
|
||||
class DatasheetFetchResult:
|
||||
def __init__(
|
||||
self,
|
||||
mpn: str,
|
||||
pdf_bytes: bytes | None = None,
|
||||
error: str | None = None,
|
||||
url: str | None = None,
|
||||
catalog_mpn: str | None = None,
|
||||
):
|
||||
self.mpn = mpn
|
||||
self.pdf_bytes = pdf_bytes
|
||||
self.error = error
|
||||
self.url = url
|
||||
self.catalog_mpn = catalog_mpn
|
||||
|
||||
@property
|
||||
def ok(self) -> bool:
|
||||
return self.pdf_bytes is not None
|
||||
|
||||
|
||||
async def fetch_datasheet(mpn: str) -> DatasheetFetchResult:
|
||||
if not settings.use_digikey:
|
||||
return DatasheetFetchResult(mpn, error="DigiKey API not configured")
|
||||
try:
|
||||
url, catalog_mpn = await _search_mpn(mpn)
|
||||
except httpx.HTTPStatusError as e:
|
||||
return DatasheetFetchResult(mpn, error=f"DigiKey search failed ({e.response.status_code})")
|
||||
except Exception as e:
|
||||
return DatasheetFetchResult(mpn, error=f"DigiKey search error: {str(e) or type(e).__name__}")
|
||||
if not url:
|
||||
return DatasheetFetchResult(mpn, error="No datasheet found on DigiKey")
|
||||
try:
|
||||
pdf_bytes = await _download_pdf(url)
|
||||
except httpx.HTTPStatusError as e:
|
||||
return DatasheetFetchResult(
|
||||
mpn, error=f"Download blocked ({e.response.status_code})", url=url, catalog_mpn=catalog_mpn,
|
||||
)
|
||||
except ValueError as e:
|
||||
return DatasheetFetchResult(mpn, error=str(e), url=url, catalog_mpn=catalog_mpn)
|
||||
except httpx.TimeoutException:
|
||||
return DatasheetFetchResult(mpn, error="Download timed out", url=url, catalog_mpn=catalog_mpn)
|
||||
except Exception as e:
|
||||
return DatasheetFetchResult(
|
||||
mpn, error=f"Download failed: {str(e) or type(e).__name__}", url=url, catalog_mpn=catalog_mpn,
|
||||
)
|
||||
return DatasheetFetchResult(mpn, pdf_bytes=pdf_bytes, url=url, catalog_mpn=catalog_mpn)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ProductParams:
|
||||
mpn: str
|
||||
parameters: list[dict[str, str]] = field(default_factory=list)
|
||||
category: str = ""
|
||||
description: str = ""
|
||||
|
||||
|
||||
class ParamsFetchResult:
|
||||
def __init__(self, mpn: str, params: ProductParams | None = None, error: str | None = None):
|
||||
self.mpn = mpn
|
||||
self.params = params
|
||||
self.error = error
|
||||
|
||||
@property
|
||||
def ok(self) -> bool:
|
||||
return self.params is not None
|
||||
|
||||
|
||||
def _parse_product_params(mpn: str, product: dict) -> ProductParams:
|
||||
parameters = []
|
||||
for p in product.get("Parameters") or product.get("parameters") or []:
|
||||
name = p.get("ParameterText") or p.get("parameterText") or ""
|
||||
value = p.get("ValueText") or p.get("valueText") or ""
|
||||
if name and value and value != "-":
|
||||
parameters.append({"name": name, "value": value})
|
||||
cat = product.get("Category") or product.get("category") or {}
|
||||
desc_obj = product.get("Description") or product.get("description") or {}
|
||||
if isinstance(desc_obj, str):
|
||||
description = desc_obj
|
||||
else:
|
||||
description = (
|
||||
desc_obj.get("ProductDescription")
|
||||
or desc_obj.get("productDescription")
|
||||
or desc_obj.get("DetailedDescription")
|
||||
or desc_obj.get("detailedDescription")
|
||||
or ""
|
||||
)
|
||||
return ProductParams(
|
||||
mpn=mpn,
|
||||
parameters=parameters,
|
||||
category=cat.get("Name") or cat.get("name") or "",
|
||||
description=description,
|
||||
)
|
||||
|
||||
|
||||
async def fetch_params(mpn: str) -> ParamsFetchResult:
|
||||
if not settings.use_digikey:
|
||||
return ParamsFetchResult(mpn, error="DigiKey API not configured")
|
||||
products: list[dict] = []
|
||||
try:
|
||||
tried: set[str] = set()
|
||||
for keyword in mpn_query_variants(mpn):
|
||||
key = keyword.upper()
|
||||
if key in tried:
|
||||
continue
|
||||
tried.add(key)
|
||||
products = await _keyword_search(keyword)
|
||||
if _find_product(mpn, products):
|
||||
break
|
||||
except httpx.HTTPStatusError as e:
|
||||
return ParamsFetchResult(mpn, error=f"DigiKey search failed ({e.response.status_code})")
|
||||
except Exception as e:
|
||||
return ParamsFetchResult(mpn, error=f"DigiKey search error: {str(e) or type(e).__name__}")
|
||||
product = _find_product(mpn, products)
|
||||
if not product:
|
||||
return ParamsFetchResult(mpn, error="No results found on DigiKey")
|
||||
params = _parse_product_params(mpn, product)
|
||||
if not params.parameters:
|
||||
return ParamsFetchResult(mpn, error="No parameters available on DigiKey")
|
||||
return ParamsFetchResult(mpn, params=params)
|
||||
@@ -0,0 +1,117 @@
|
||||
"""Per-event JSON objects under users/.../events/ for SSE across processes."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from datetime import datetime, timezone
|
||||
from typing import AsyncIterator
|
||||
|
||||
from backend.services.projects import project_prefix
|
||||
from backend.services.storage import StorageBackend
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_SEQ_WIDTH = 10
|
||||
_FILENAME_FMT = f"{{seq:0{_SEQ_WIDTH}d}}.json"
|
||||
|
||||
TERMINAL_EVENTS = frozenset({
|
||||
"pipeline_complete",
|
||||
"pipeline_error",
|
||||
"pipeline_cancelled",
|
||||
"pipeline_paused",
|
||||
"placement_complete",
|
||||
"placement_error",
|
||||
"placement_cancelled",
|
||||
"pcb_complete",
|
||||
"pcb_error",
|
||||
"pcb_cancelled",
|
||||
})
|
||||
|
||||
|
||||
def _events_prefix(user_id: str, project_id: str) -> str:
|
||||
return f"{project_prefix(user_id, project_id)}/events/"
|
||||
|
||||
|
||||
def _seq_from_key(key: str) -> int | None:
|
||||
name = key.rsplit("/", 1)[-1]
|
||||
if not name.endswith(".json"):
|
||||
return None
|
||||
try:
|
||||
return int(name[:-5])
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
|
||||
class GCSEventBroker:
|
||||
def __init__(self, storage: StorageBackend, user_id: str) -> None:
|
||||
self.storage = storage
|
||||
self.user_id = user_id
|
||||
self._seq: dict[str, int] = {}
|
||||
|
||||
def subscribe(self, project_id: str) -> asyncio.Queue:
|
||||
raise NotImplementedError("use event_bridge.tail_events instead")
|
||||
|
||||
def unsubscribe(self, project_id: str, q: asyncio.Queue) -> None:
|
||||
return
|
||||
|
||||
def clear_history(self, project_id: str) -> None:
|
||||
prefix = _events_prefix(self.user_id, project_id)
|
||||
try:
|
||||
self.storage.delete_prefix(prefix)
|
||||
except Exception:
|
||||
logger.exception("failed to clear events at %s", prefix)
|
||||
self._seq[project_id] = 0
|
||||
|
||||
def publish(self, project_id: str, event: str, data: dict) -> None:
|
||||
seq = self._seq.get(project_id, 0)
|
||||
self._seq[project_id] = seq + 1
|
||||
key = _events_prefix(self.user_id, project_id) + _FILENAME_FMT.format(seq=seq)
|
||||
try:
|
||||
self.storage.write_json(key, {
|
||||
"seq": seq,
|
||||
"ts": datetime.now(timezone.utc).isoformat(),
|
||||
"event": event,
|
||||
"data": data,
|
||||
})
|
||||
except Exception:
|
||||
logger.exception("failed to write event %s to %s", event, key)
|
||||
|
||||
|
||||
async def tail_events(
|
||||
storage: StorageBackend,
|
||||
user_id: str,
|
||||
project_id: str,
|
||||
*,
|
||||
poll_interval: float = 0.5,
|
||||
heartbeat_interval: float = 15.0,
|
||||
terminal_events: frozenset[str] | None = None,
|
||||
) -> AsyncIterator[dict]:
|
||||
stop_on = terminal_events if terminal_events is not None else TERMINAL_EVENTS
|
||||
prefix = _events_prefix(user_id, project_id)
|
||||
last_seen_key: str | None = None
|
||||
last_emit_ts = 0.0
|
||||
while True:
|
||||
try:
|
||||
keys = storage.list_prefix_after(prefix, after_key=last_seen_key)
|
||||
except Exception:
|
||||
logger.exception("event tail list failed for %s", prefix)
|
||||
keys = []
|
||||
emitted_any = False
|
||||
for key in keys:
|
||||
try:
|
||||
msg = storage.read_json(key)
|
||||
except Exception:
|
||||
logger.exception("event tail read failed for %s", key)
|
||||
continue
|
||||
yield msg
|
||||
emitted_any = True
|
||||
last_seen_key = key
|
||||
last_emit_ts = asyncio.get_event_loop().time()
|
||||
if msg.get("event") in stop_on:
|
||||
return
|
||||
now = asyncio.get_event_loop().time()
|
||||
if not emitted_any and now - last_emit_ts >= heartbeat_interval:
|
||||
yield {"event": "heartbeat", "data": {}}
|
||||
last_emit_ts = now
|
||||
await asyncio.sleep(poll_interval)
|
||||
@@ -0,0 +1,310 @@
|
||||
"""Dispatch pipeline workers: local subprocess in self-host, Cloud Run when GCS is set."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
from pathlib import Path
|
||||
from typing import Literal
|
||||
|
||||
from backend.config import settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
ExecutionState = Literal[
|
||||
"pending", "running", "succeeded", "failed", "cancelled", "unknown"
|
||||
]
|
||||
|
||||
_local_procs: dict[str, subprocess.Popen] = {}
|
||||
_local_procs_lock = threading.Lock()
|
||||
|
||||
|
||||
def _local_execution_name(project_id: str) -> str:
|
||||
return f"local/projects/{project_id}"
|
||||
|
||||
|
||||
def _pid_path(project_id: str) -> Path:
|
||||
return settings.data_dir / "workers" / f"{project_id}.pid"
|
||||
|
||||
|
||||
def _write_pid(project_id: str, pid: int) -> None:
|
||||
path = _pid_path(project_id)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(str(pid))
|
||||
|
||||
|
||||
def _pid_alive(project_id: str) -> bool | None:
|
||||
path = _pid_path(project_id)
|
||||
if not path.is_file():
|
||||
return None
|
||||
try:
|
||||
pid = int(path.read_text().strip())
|
||||
except ValueError:
|
||||
return False
|
||||
try:
|
||||
os.kill(pid, 0)
|
||||
except OSError:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _spawn_local_subprocess(
|
||||
project_id: str,
|
||||
user_id: str,
|
||||
*,
|
||||
resume: bool,
|
||||
free: bool,
|
||||
mode: str = "run",
|
||||
regen_stages: list[str] | None = None,
|
||||
proc_key: str | None = None,
|
||||
execution_name: str | None = None,
|
||||
) -> str:
|
||||
key = proc_key or project_id
|
||||
name = execution_name or _local_execution_name(project_id)
|
||||
env = os.environ.copy()
|
||||
env["PROJECT_ID"] = project_id
|
||||
env["USER_ID"] = user_id
|
||||
env["RESUME"] = "1" if resume else "0"
|
||||
env["FREE"] = "1" if free else "0"
|
||||
env["MODE"] = mode
|
||||
if regen_stages:
|
||||
env["REGEN_STAGES"] = ",".join(regen_stages)
|
||||
env["EXECUTION_NAME"] = name
|
||||
proc = subprocess.Popen(
|
||||
[sys.executable, "-m", "backend.pipeline_worker"],
|
||||
env=env,
|
||||
stdin=subprocess.DEVNULL,
|
||||
)
|
||||
_write_pid(key, proc.pid)
|
||||
with _local_procs_lock:
|
||||
prior = _local_procs.pop(key, None)
|
||||
if prior is not None:
|
||||
try:
|
||||
prior.terminate()
|
||||
except Exception:
|
||||
pass
|
||||
_local_procs[key] = proc
|
||||
logger.info("spawned worker pid=%s project=%s mode=%s", proc.pid, project_id, mode)
|
||||
return name
|
||||
|
||||
|
||||
def _local_state(project_id: str) -> ExecutionState:
|
||||
with _local_procs_lock:
|
||||
proc = _local_procs.get(project_id)
|
||||
if proc is not None:
|
||||
rc = proc.poll()
|
||||
if rc is None:
|
||||
return "running"
|
||||
if rc == 0:
|
||||
return "succeeded"
|
||||
if rc < 0:
|
||||
return "cancelled"
|
||||
return "failed"
|
||||
alive = _pid_alive(project_id)
|
||||
if alive is True:
|
||||
return "running"
|
||||
if alive is False:
|
||||
return "failed"
|
||||
return "unknown"
|
||||
|
||||
|
||||
def _local_cancel(project_id: str) -> None:
|
||||
with _local_procs_lock:
|
||||
proc = _local_procs.get(project_id)
|
||||
if proc is None or proc.poll() is not None:
|
||||
return
|
||||
try:
|
||||
proc.terminate()
|
||||
except Exception:
|
||||
logger.exception("failed to terminate worker for %s", project_id)
|
||||
|
||||
|
||||
def _gcp_project() -> str:
|
||||
if settings.pipeline_worker_project:
|
||||
return settings.pipeline_worker_project
|
||||
proj = os.environ.get("GOOGLE_CLOUD_PROJECT") or os.environ.get("GCLOUD_PROJECT")
|
||||
if proj:
|
||||
return proj
|
||||
try:
|
||||
import requests
|
||||
|
||||
resp = requests.get(
|
||||
"http://metadata.google.internal/computeMetadata/v1/project/project-id",
|
||||
headers={"Metadata-Flavor": "Google"},
|
||||
timeout=2.0,
|
||||
)
|
||||
if resp.ok:
|
||||
return resp.text.strip()
|
||||
except Exception:
|
||||
pass
|
||||
raise RuntimeError(
|
||||
"Could not resolve GCP project. Set PIPELINE_WORKER_PROJECT or GOOGLE_CLOUD_PROJECT."
|
||||
)
|
||||
|
||||
|
||||
def _job_resource_name() -> str:
|
||||
return (
|
||||
f"projects/{_gcp_project()}/locations/{settings.pipeline_worker_region}"
|
||||
f"/jobs/{settings.pipeline_worker_job_name}"
|
||||
)
|
||||
|
||||
|
||||
def _jobs_client():
|
||||
from google.cloud import run_v2
|
||||
|
||||
return run_v2.JobsClient()
|
||||
|
||||
|
||||
def _executions_client():
|
||||
from google.cloud import run_v2
|
||||
|
||||
return run_v2.ExecutionsClient()
|
||||
|
||||
|
||||
def _enqueue_cloud_run_job(
|
||||
project_id: str,
|
||||
user_id: str,
|
||||
*,
|
||||
resume: bool,
|
||||
free: bool,
|
||||
mode: str = "run",
|
||||
regen_stages: list[str] | None = None,
|
||||
) -> str:
|
||||
from google.cloud import run_v2
|
||||
|
||||
env_overrides = [
|
||||
run_v2.EnvVar(name="PROJECT_ID", value=project_id),
|
||||
run_v2.EnvVar(name="USER_ID", value=user_id),
|
||||
run_v2.EnvVar(name="RESUME", value="1" if resume else "0"),
|
||||
run_v2.EnvVar(name="FREE", value="1" if free else "0"),
|
||||
run_v2.EnvVar(name="MODE", value=mode),
|
||||
]
|
||||
if regen_stages:
|
||||
env_overrides.append(run_v2.EnvVar(name="REGEN_STAGES", value=",".join(regen_stages)))
|
||||
overrides = run_v2.RunJobRequest.Overrides(
|
||||
container_overrides=[
|
||||
run_v2.RunJobRequest.Overrides.ContainerOverride(env=env_overrides),
|
||||
],
|
||||
)
|
||||
operation = _jobs_client().run_job(
|
||||
request=run_v2.RunJobRequest(name=_job_resource_name(), overrides=overrides)
|
||||
)
|
||||
metadata = operation.metadata
|
||||
name = getattr(metadata, "name", None) if metadata is not None else None
|
||||
if not name:
|
||||
name = operation.operation.name # type: ignore[union-attr]
|
||||
if not name:
|
||||
raise RuntimeError("Cloud Run RunJob returned no execution name")
|
||||
logger.info("enqueued Cloud Run execution %s for %s", name, project_id)
|
||||
return name
|
||||
|
||||
|
||||
def _cloud_run_state(execution_name: str) -> ExecutionState:
|
||||
try:
|
||||
ex = _executions_client().get_execution(name=execution_name)
|
||||
except Exception:
|
||||
logger.exception("get_execution failed for %s", execution_name)
|
||||
return "unknown"
|
||||
if ex.completion_time is None or ex.completion_time.seconds == 0:
|
||||
if ex.start_time and ex.start_time.seconds:
|
||||
return "running"
|
||||
return "pending"
|
||||
failed = int(getattr(ex, "failed_count", 0) or 0)
|
||||
cancelled = int(getattr(ex, "cancelled_count", 0) or 0)
|
||||
succeeded = int(getattr(ex, "succeeded_count", 0) or 0)
|
||||
if cancelled > 0 and succeeded == 0:
|
||||
return "cancelled"
|
||||
if failed > 0:
|
||||
return "failed"
|
||||
if succeeded > 0:
|
||||
return "succeeded"
|
||||
return "unknown"
|
||||
|
||||
|
||||
def _cloud_run_cancel(execution_name: str) -> None:
|
||||
try:
|
||||
from google.cloud import run_v2
|
||||
|
||||
_executions_client().cancel_execution(
|
||||
request=run_v2.CancelExecutionRequest(name=execution_name)
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("cancel_execution failed for %s", execution_name)
|
||||
|
||||
|
||||
def use_cloud_run_jobs() -> bool:
|
||||
return bool(settings.gcs_bucket)
|
||||
|
||||
|
||||
def enqueue_pipeline(
|
||||
project_id: str, user_id: str, *, resume: bool = False, free: bool = False,
|
||||
) -> str:
|
||||
if use_cloud_run_jobs():
|
||||
return _enqueue_cloud_run_job(project_id, user_id, resume=resume, free=free)
|
||||
return _spawn_local_subprocess(project_id, user_id, resume=resume, free=free)
|
||||
|
||||
|
||||
def enqueue_pipeline_regen(project_id: str, user_id: str, *, stages: list[str]) -> str:
|
||||
if not stages:
|
||||
raise ValueError("regen requires at least one stage")
|
||||
if use_cloud_run_jobs():
|
||||
return _enqueue_cloud_run_job(
|
||||
project_id, user_id, resume=False, free=True, mode="regen", regen_stages=stages,
|
||||
)
|
||||
return _spawn_local_subprocess(
|
||||
project_id, user_id, resume=False, free=True, mode="regen", regen_stages=stages,
|
||||
)
|
||||
|
||||
|
||||
def get_execution_state(execution_name: str | None) -> ExecutionState:
|
||||
if not execution_name:
|
||||
return "unknown"
|
||||
if execution_name.startswith("local/projects/"):
|
||||
return _local_state(execution_name.split("/", 2)[-1])
|
||||
if execution_name.startswith("local/placement/"):
|
||||
return _local_state(f"placement:{execution_name.split('/', 2)[-1]}")
|
||||
if execution_name.startswith("local/pcb/"):
|
||||
return _local_state(f"pcb:{execution_name.split('/', 2)[-1]}")
|
||||
return _cloud_run_state(execution_name)
|
||||
|
||||
|
||||
def cancel_execution(execution_name: str | None) -> None:
|
||||
if not execution_name:
|
||||
return
|
||||
if execution_name.startswith("local/projects/"):
|
||||
_local_cancel(execution_name.split("/", 2)[-1])
|
||||
return
|
||||
if execution_name.startswith("local/placement/"):
|
||||
_local_cancel(f"placement:{execution_name.split('/', 2)[-1]}")
|
||||
return
|
||||
if execution_name.startswith("local/pcb/"):
|
||||
_local_cancel(f"pcb:{execution_name.split('/', 2)[-1]}")
|
||||
return
|
||||
_cloud_run_cancel(execution_name)
|
||||
|
||||
|
||||
def enqueue_placement_pipeline(project_id: str, user_id: str) -> str:
|
||||
if use_cloud_run_jobs():
|
||||
return _enqueue_cloud_run_job(
|
||||
project_id, user_id, resume=False, free=True, mode="placement",
|
||||
)
|
||||
return _spawn_local_subprocess(
|
||||
project_id, user_id, resume=False, free=True, mode="placement",
|
||||
proc_key=f"placement:{project_id}",
|
||||
execution_name=f"local/placement/{project_id}",
|
||||
)
|
||||
|
||||
|
||||
def enqueue_pcb_pipeline(project_id: str, user_id: str) -> str:
|
||||
if use_cloud_run_jobs():
|
||||
return _enqueue_cloud_run_job(
|
||||
project_id, user_id, resume=False, free=False, mode="pcb",
|
||||
)
|
||||
return _spawn_local_subprocess(
|
||||
project_id, user_id, resume=False, free=False, mode="pcb",
|
||||
proc_key=f"pcb:{project_id}",
|
||||
execution_name=f"local/pcb/{project_id}",
|
||||
)
|
||||
@@ -0,0 +1,40 @@
|
||||
"""Leftover dispatcher/cost/DigiKey modules load from src."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import backend.services.cost_estimator as cost_estimator
|
||||
import backend.services.digikey as digikey
|
||||
import backend.services.event_bridge as event_bridge
|
||||
import backend.services.job_runner as job_runner
|
||||
from backend.services.digikey import _find_product
|
||||
from backend.services.job_runner import get_execution_state
|
||||
|
||||
|
||||
def test_dispatcher_modules_are_src():
|
||||
for mod, name in (
|
||||
(job_runner, "job_runner.py"),
|
||||
(event_bridge, "event_bridge.py"),
|
||||
(cost_estimator, "cost_estimator.py"),
|
||||
(digikey, "digikey.py"),
|
||||
):
|
||||
path = Path(mod.__file__).resolve()
|
||||
assert path.name == name
|
||||
assert "src" in path.parts
|
||||
assert "Native Periscope overlay" not in path.read_text(encoding="utf-8")[:400]
|
||||
|
||||
|
||||
def test_digikey_rejects_unrelated_and_accepts_family():
|
||||
products = [{"ManufacturerProductNumber": "CH340E"}]
|
||||
assert _find_product("CH340", products) is None
|
||||
products = [{"ManufacturerProductNumber": "24AA025E64-I/SN"}]
|
||||
assert _find_product("24AA025E64", products) is not None
|
||||
|
||||
|
||||
def test_stale_pid_is_failed(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(job_runner.settings, "data_dir", tmp_path)
|
||||
pid_path = tmp_path / "workers" / "proj.pid"
|
||||
pid_path.parent.mkdir(parents=True)
|
||||
pid_path.write_text("99999999")
|
||||
assert get_execution_state("local/projects/proj") == "failed"
|
||||
Reference in New Issue
Block a user