diff --git a/periscope/src/backend/services/cost_estimator.py b/periscope/src/backend/services/cost_estimator.py new file mode 100644 index 0000000..5392e53 --- /dev/null +++ b/periscope/src/backend/services/cost_estimator.py @@ -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, + ) diff --git a/periscope/src/backend/services/digikey.py b/periscope/src/backend/services/digikey.py new file mode 100644 index 0000000..1d473a6 --- /dev/null +++ b/periscope/src/backend/services/digikey.py @@ -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) diff --git a/periscope/src/backend/services/event_bridge.py b/periscope/src/backend/services/event_bridge.py new file mode 100644 index 0000000..3d1b53d --- /dev/null +++ b/periscope/src/backend/services/event_bridge.py @@ -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) diff --git a/periscope/src/backend/services/job_runner.py b/periscope/src/backend/services/job_runner.py new file mode 100644 index 0000000..fe431cc --- /dev/null +++ b/periscope/src/backend/services/job_runner.py @@ -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}", + ) diff --git a/tests/test_periscope_services_dispatch_rewrite.py b/tests/test_periscope_services_dispatch_rewrite.py new file mode 100644 index 0000000..e2ffb33 --- /dev/null +++ b/tests/test_periscope_services_dispatch_rewrite.py @@ -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"