Files
planet/backend/app/services/compute_center_locations.py
rayd1o 5bf5c73ca0
Some checks failed
ci / backend (push) Has been cancelled
ci / frontend (push) Has been cancelled
ci / delivery (push) Has been cancelled
release / images (push) Has been cancelled
release: bump version to 0.66.0
2026-05-26 03:41:47 +08:00

887 lines
30 KiB
Python

"""Compute-center location resolver, built on the shared location pipeline.
This module is a thin domain wrapper that wires up
:mod:`app.services.location` for compute centers:
SourceCoordinates
The online Nominatim step is intentionally reserved for the user-triggered
``collect-location`` flow. The regular GeoJSON endpoint runs during Earth
startup, so it must stay local and deterministic.
For the full design and the reason behind the abstraction (compute centers,
BGP collectors, BGP events, and future entities all share one pipeline),
see ``docs/technical/zh/location-pipeline-development.md``.
The ``ComputeCenterLocation`` dataclass and the public function signatures are
preserved verbatim so existing callers and tests do not need to change.
"""
from __future__ import annotations
from dataclasses import dataclass
from datetime import UTC, datetime
from functools import lru_cache
from typing import Any
import httpx
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from app.core.collected_data_fields import get_record_field
from app.models.collected_data import CollectedData
from app.models.compute_center_location import ComputeCenterLocationRecord
from app.services.location import (
LocationCandidate,
LocationPipeline,
LocationQuery,
NominatimResolver,
ResolverOutput,
SourceCoordinatesResolver,
build_default_nominatim_geocoder,
coerce_str,
normalize_country_text,
normalize_text,
parse_float,
)
ROR_SEARCH_URL = "https://api.ror.org/v2/organizations"
DEFAULT_ROR_USER_AGENT = "planet-earth-location-resolver/1.0"
DEFAULT_ROR_TIMEOUT_SECONDS = 8.0
RENDERABLE_PRECISIONS: tuple[str, ...] = ("precise", "site", "city")
FORBIDDEN_PRECISIONS: tuple[str, ...] = (
"country",
"estimated_country",
"country_major_compute_city",
"region",
"unknown",
)
# ── Public dataclasses ──────────────────────────────────────────────
@dataclass(frozen=True)
class ComputeCenterLocation:
latitude: float | None
longitude: float | None
location_precision: str
geography_mode: str
is_estimated: bool
estimated_reason: str | None = None
location_confidence: float | None = None
location_source: str | None = None
location_source_note: str | None = None
location_verified_at: str | None = None
matched_location_name: str | None = None
needs_confirmation: bool = False
city: str | None = None
region: str | None = None
country: str | None = None
@property
def is_renderable(self) -> bool:
if self.latitude in (None, 0.0) or self.longitude in (None, 0.0):
return False
return self.location_precision in RENDERABLE_PRECISIONS
def to_geojson_properties(self) -> dict[str, Any]:
return {
"latitude": self.latitude,
"longitude": self.longitude,
"location_precision": self.location_precision,
"geography_mode": self.geography_mode,
"is_estimated": self.is_estimated,
"estimated_reason": self.estimated_reason,
"location_confidence": self.location_confidence,
"location_source": self.location_source,
"location_source_note": self.location_source_note,
"location_verified_at": self.location_verified_at,
"matched_location_name": self.matched_location_name,
"needs_confirmation": self.needs_confirmation,
}
@dataclass(frozen=True)
class ResolutionDiagnostic:
failure_reason: str
attempted_queries: tuple[str, ...] = ()
record_id: int | None = None
source: str | None = None
source_id: str | None = None
name: str | None = None
country: str | None = None
city: str | None = None
site: str | None = None
operator: str | None = None
def to_dict(self) -> dict[str, Any]:
return {
"failure_reason": self.failure_reason,
"attempted_queries": list(self.attempted_queries),
"record_id": self.record_id,
"source": self.source,
"source_id": self.source_id,
"name": self.name,
"country": self.country,
"city": self.city,
"site": self.site,
"operator": self.operator,
}
@dataclass(frozen=True)
class ResolutionResult:
location: ComputeCenterLocation | None
diagnostic: ResolutionDiagnostic | None
@property
def is_resolved(self) -> bool:
return bool(self.location and self.location.is_renderable)
# ── Geocoder (kept at module level so tests can monkeypatch + cache_clear) ──
_geocode_online = build_default_nominatim_geocoder()
# ── Stored location cache ───────────────────────────────────────────
COMPUTE_CENTER_LOCATION_CACHE: dict[str, dict[str, Any]] = {}
def _cache_key(source: str | None, source_id: str | None) -> str:
return f"{coerce_str(source)}:{coerce_str(source_id)}"
def set_compute_center_location_cache(
locations: dict[str, dict[str, Any]],
) -> None:
COMPUTE_CENTER_LOCATION_CACHE.clear()
COMPUTE_CENTER_LOCATION_CACHE.update(
{coerce_str(key): dict(value) for key, value in locations.items()}
)
async def refresh_compute_center_location_cache(
session: AsyncSession,
) -> dict[str, dict[str, Any]]:
result = await session.execute(select(ComputeCenterLocationRecord))
records = result.scalars().all()
cache = {}
for record in records:
if not hasattr(record, "to_location_dict"):
continue
if not record.source or not record.source_id:
continue
cache[_cache_key(record.source, record.source_id)] = record.to_location_dict()
set_compute_center_location_cache(cache)
return cache
def get_compute_center_location_dict(
source: str | None,
source_id: str | None,
) -> dict[str, Any]:
return dict(COMPUTE_CENTER_LOCATION_CACHE.get(_cache_key(source, source_id), {}))
# ── Pipeline construction ──────────────────────────────────────────
@lru_cache(maxsize=512)
def _lookup_ror_organization(query: str) -> dict[str, Any] | None:
"""Lookup a research organization in ROR for user-triggered candidates."""
if not query:
return None
response = httpx.get(
ROR_SEARCH_URL,
params={"query": query},
headers={"User-Agent": DEFAULT_ROR_USER_AGENT},
timeout=DEFAULT_ROR_TIMEOUT_SECONDS,
)
response.raise_for_status()
payload = response.json()
items = payload.get("items") if isinstance(payload, dict) else None
if not isinstance(items, list) or not items:
return None
first = items[0]
if not isinstance(first, dict):
return None
organization = first.get("organization")
if isinstance(organization, dict):
return organization
return first
def _compute_center_ror_query_plan(
query: LocationQuery,
) -> list[tuple[str, tuple[str, ...]]]:
extra = query.extra or {}
raw_parts: list[tuple[str, str]] = [
("site", coerce_str(extra.get("site"))),
("operator", coerce_str(extra.get("operator"))),
("organization", coerce_str(extra.get("organization"))),
]
for field, value in tuple(raw_parts):
if "/" not in value:
continue
raw_parts.extend(
(field, part.strip())
for part in value.split("/")
if len(part.strip()) >= 3
)
plan: list[tuple[str, tuple[str, ...]]] = []
seen: set[str] = set()
for field, value in raw_parts:
key = normalize_text(value)
if not key or key in seen:
continue
seen.add(key)
plan.append((value, (field,)))
return plan
def _organization_label(organization: dict[str, Any], fallback: str) -> str:
names = organization.get("names")
if isinstance(names, list):
for name in names:
if not isinstance(name, dict):
continue
types = name.get("types")
if isinstance(types, list) and "ror_display" in types:
value = coerce_str(name.get("value"))
if value:
return value
for name in names:
if isinstance(name, dict):
value = coerce_str(name.get("value"))
if value:
return value
return fallback
class ROROrganizationResolver:
"""Resolve source-provided organization/site text through the open ROR API."""
name = "ror_organization_registry"
def __init__(
self,
*,
query_plan_builder=_compute_center_ror_query_plan,
lookup=lambda q: _lookup_ror_organization(q),
confidence: float = 0.68,
) -> None:
self._query_plan_builder = query_plan_builder
self._lookup = lookup
self._confidence = confidence
def resolve(self, query: LocationQuery):
from app.services.location import ResolverOutput
from app.services.location.text import parse_float
attempted: list[str] = []
candidates: list[LocationCandidate] = []
context_country = normalize_text(normalize_country_text(query.country))
for ror_query, matched_fields in self._query_plan_builder(query):
attempted.append(f"ror:{ror_query}")
try:
organization = self._lookup(ror_query)
except Exception:
continue
if not isinstance(organization, dict):
continue
locations = organization.get("locations")
if not isinstance(locations, list) or not locations:
continue
location = locations[0]
if not isinstance(location, dict):
continue
details = location.get("geonames_details")
if not isinstance(details, dict):
continue
latitude = parse_float(details.get("lat"))
longitude = parse_float(details.get("lng"))
if latitude in (None, 0.0) or longitude in (None, 0.0):
continue
country = normalize_country_text(details.get("country_name"))
if context_country and normalize_text(country) != context_country:
continue
city = coerce_str(details.get("name")) or None
region = coerce_str(details.get("country_subdivision_name")) or None
display_name = _organization_label(organization, ror_query)
ror_id = coerce_str(organization.get("id"))
geonames_id = location.get("geonames_id")
source_note = (
f"ROR organization match: {display_name}"
+ (f" ({ror_id})" if ror_id else "")
+ (f"; GeoNames {geonames_id}" if geonames_id else "")
)
candidates.append(
LocationCandidate(
latitude=latitude,
longitude=longitude,
display_name=display_name,
precision="city",
confidence=self._confidence,
query=ror_query,
source=self.name,
source_note=source_note,
matched_fields=matched_fields,
needs_confirmation=True,
city=city,
region=region,
country=country or query.country,
matched_location_name=display_name,
location_verified_at=None,
suggested_registry_entry=None,
)
)
return ResolverOutput(
candidates=tuple(candidates),
attempted_queries=tuple(attempted),
)
class StoredComputeCenterLocationResolver:
"""Resolve a compute center through the DB-backed current-location cache."""
name = "stored_compute_center_location"
def resolve(self, query: LocationQuery) -> ResolverOutput:
extra = query.extra or {}
stored = get_compute_center_location_dict(
coerce_str(extra.get("source")),
coerce_str(extra.get("source_id")),
)
if not stored:
return ResolverOutput()
latitude = parse_float(stored.get("latitude"))
longitude = parse_float(stored.get("longitude"))
if latitude in (None, 0.0) or longitude in (None, 0.0):
return ResolverOutput()
return ResolverOutput(
candidates=(
LocationCandidate(
latitude=latitude,
longitude=longitude,
display_name=stored.get("name") or query.name or "Compute center",
precision=stored.get("precision") or "city",
confidence=float(stored.get("confidence") or 0.85),
query=f"stored_compute_center_location::{stored.get('source')}:{stored.get('source_id')}",
source=self.name,
source_note=stored.get("source_note"),
matched_fields=("source", "source_id"),
needs_confirmation=bool(stored.get("needs_confirmation")),
city=stored.get("city") or query.city,
region=None,
country=stored.get("country") or query.country,
matched_location_name=stored.get("site") or stored.get("name") or query.name,
location_verified_at=stored.get("verified_at"),
suggested_registry_entry=None,
),
)
)
def _short_system_name(name: Any) -> str:
"""Strip vendor/system suffix from TOP500 names like ``"El Capitan - HPE Cray ..."``."""
text = coerce_str(name)
if not text:
return ""
head = text.split(" - ", 1)[0].strip()
return head or text
def _record_context(record: Any, metadata: dict[str, Any]) -> dict[str, str]:
name = coerce_str(getattr(record, "name", None))
return {
"source": coerce_str(getattr(record, "source", None)),
"source_id": coerce_str(getattr(record, "source_id", None)),
"name": name,
"name_short": _short_system_name(name),
"city": coerce_str(get_record_field(record, "city")),
"country": coerce_str(get_record_field(record, "country")),
"site": coerce_str(metadata.get("site") or metadata.get("organization")),
"operator": coerce_str(
metadata.get("operator")
or metadata.get("organization")
or metadata.get("owner")
or metadata.get("manufacturer")
),
"organization": coerce_str(metadata.get("organization")),
}
def _context_to_query(
context: dict[str, str],
*,
source_lat: float | None = None,
source_lon: float | None = None,
) -> LocationQuery:
name = context.get("name") or None
name_short = context.get("name_short") or ""
aliases: tuple[str, ...] = ()
if name_short and name_short != name:
aliases = (name_short,)
return LocationQuery(
name=name,
aliases=aliases,
city=context.get("city") or None,
country=context.get("country") or None,
source_latitude=source_lat,
source_longitude=source_lon,
extra={
"source": context.get("source") or "",
"source_id": context.get("source_id") or "",
"site": context.get("site") or "",
"operator": context.get("operator") or "",
"organization": context.get("organization") or "",
},
)
def _compute_center_query_plan(
query: LocationQuery,
) -> list[tuple[str, tuple[str, ...]]]:
"""Build the Nominatim query plan for a compute-center query.
Mirrors the legacy ``_build_online_query_plan`` ordering exactly.
"""
name = query.name or ""
name_short = (query.aliases[0] if query.aliases else "") or name
extra = query.extra or {}
site = str(extra.get("site") or "")
operator = str(extra.get("operator") or "")
city = query.city or ""
country = query.country or ""
plan: list[tuple[str, tuple[str, ...]]] = []
def add(parts: list[tuple[str, str]]) -> None:
non_empty = [(field, value) for field, value in parts if value]
if not non_empty:
return
seen: set[str] = set()
cleaned: list[str] = []
fields: list[str] = []
for field, value in non_empty:
key = normalize_text(value)
if not key or key in seen:
continue
seen.add(key)
cleaned.append(value)
fields.append(field)
if not cleaned:
return
composed = ", ".join(cleaned)
if not any(composed == existing for existing, _ in plan):
plan.append((composed, tuple(fields)))
add([("site", site), ("country", country)])
add([("operator", operator), ("city", city), ("country", country)])
add([("name", name_short), ("operator", operator), ("country", country)])
add([("name", name_short), ("site", site)])
add([("name", name_short), ("country", country)])
add([("name", name_short), ("city", city), ("country", country)])
add([("city", city), ("country", country)])
if name and name != name_short:
add([("name", name), ("country", country)])
return plan
COMPUTE_CENTER_PIPELINE = LocationPipeline(
[
SourceCoordinatesResolver(),
StoredComputeCenterLocationResolver(),
],
failure_reason=(
"Could not resolve to city-level coordinates from source coords"
" or stored compute-center location."
),
)
COMPUTE_CENTER_COLLECTION_PIPELINE = LocationPipeline(
[
SourceCoordinatesResolver(),
ROROrganizationResolver(),
NominatimResolver(
query_plan_builder=_compute_center_query_plan,
# Late-binding so test monkeypatching of ``_geocode_online`` works.
geocoder=lambda q: _geocode_online(q),
),
],
failure_reason=(
"Could not resolve to city-level coordinates from source coords"
", ROR organization lookup, or online geocoding."
),
)
# ── Candidate → ComputeCenterLocation conversion ───────────────────
_GEOGRAPHY_MODE_BY_SOURCE = {
"source_coordinates": "source_coordinates",
"stored_compute_center_location": "stored_compute_center_location",
"ror_organization_registry": "ror_organization",
"nominatim_online_geocode": "online_geocode",
}
def _candidate_to_location(
candidate: LocationCandidate,
*,
context: dict[str, str],
) -> ComputeCenterLocation:
geography_mode = _GEOGRAPHY_MODE_BY_SOURCE.get(candidate.source, "online_geocode")
is_estimated = candidate.needs_confirmation or candidate.source.startswith(
"nominatim"
)
estimated_reason: str | None
if candidate.source == "source_coordinates":
estimated_reason = None
elif candidate.source == "stored_compute_center_location":
estimated_reason = candidate.source_note
elif candidate.source == "ror_organization_registry":
fields_summary = ", ".join(candidate.matched_fields) or "organization"
estimated_reason = (
f"Resolved by ROR organization lookup '{candidate.query}' "
f"(matched fields: {fields_summary})"
)
elif candidate.source == "nominatim_online_geocode":
fields_summary = ", ".join(candidate.matched_fields) or "name"
estimated_reason = (
f"Resolved by online geocoding query '{candidate.query}' "
f"(matched fields: {fields_summary})"
)
else:
estimated_reason = candidate.source_note
country = (
candidate.country
or normalize_country_text(context.get("country"))
or context.get("country")
or None
)
return ComputeCenterLocation(
latitude=candidate.latitude,
longitude=candidate.longitude,
location_precision=candidate.precision,
geography_mode=geography_mode,
is_estimated=is_estimated,
estimated_reason=estimated_reason,
location_confidence=candidate.confidence,
location_source=candidate.source,
location_source_note=candidate.source_note,
location_verified_at=candidate.location_verified_at,
matched_location_name=candidate.matched_location_name
or context.get("name")
or None,
needs_confirmation=candidate.needs_confirmation,
city=candidate.city or context.get("city") or None,
region=candidate.region,
country=country,
)
def _diagnostic_for(
record: Any,
context: dict[str, str],
*,
failure_reason: str,
attempted_queries: tuple[str, ...] = (),
) -> ResolutionDiagnostic:
return ResolutionDiagnostic(
failure_reason=failure_reason,
attempted_queries=attempted_queries,
record_id=getattr(record, "id", None),
source=getattr(record, "source", None),
source_id=getattr(record, "source_id", None),
name=context.get("name") or getattr(record, "name", None),
country=context.get("country") or None,
city=context.get("city") or None,
site=context.get("site") or None,
operator=context.get("operator") or None,
)
# ── Public API ─────────────────────────────────────────────────────
def resolve_compute_center_location(
record: Any,
metadata: dict[str, Any] | None = None,
) -> ComputeCenterLocation:
"""Backwards-compatible thin wrapper returning the renderable location only.
Records that cannot be resolved to city-level get a placeholder
:class:`ComputeCenterLocation` with ``location_precision='unknown'``.
Callers should generally prefer :func:`resolve_compute_center_location_full`.
"""
full = resolve_compute_center_location_full(record, metadata)
return full.location or ComputeCenterLocation(
latitude=None,
longitude=None,
location_precision="unknown",
geography_mode="unresolved",
is_estimated=True,
estimated_reason="No resolvable location hints",
location_confidence=0.0,
location_source="unknown",
location_source_note=(
"No source coordinates, ROR organization match, or online"
" geocoding result."
),
matched_location_name=None,
needs_confirmation=False,
)
def resolve_compute_center_location_full(
record: Any,
metadata: dict[str, Any] | None = None,
*,
allow_online: bool = False,
) -> ResolutionResult:
metadata = metadata or {}
context = _record_context(record, metadata)
from app.services.location.text import parse_float as _parse_float
source_lat = _parse_float(get_record_field(record, "latitude"))
source_lon = _parse_float(get_record_field(record, "longitude"))
if source_lat in (None, 0.0):
source_lat = None
if source_lon in (None, 0.0):
source_lon = None
query = _context_to_query(
context, source_lat=source_lat, source_lon=source_lon
)
pipeline = (
COMPUTE_CENTER_COLLECTION_PIPELINE
if allow_online
else COMPUTE_CENTER_PIPELINE
)
pipeline_result = pipeline.resolve_best(query)
if pipeline_result.location and pipeline_result.location.precision in RENDERABLE_PRECISIONS:
location = _candidate_to_location(pipeline_result.location, context=context)
return ResolutionResult(location=location, diagnostic=None)
return ResolutionResult(
location=None,
diagnostic=_diagnostic_for(
record,
context,
failure_reason=(
"Could not resolve to city-level coordinates from source coords"
", ROR organization lookup, or online geocoding."
if allow_online
else (
"Could not resolve to city-level coordinates from source coords"
" or stored compute-center location."
)
),
attempted_queries=pipeline_result.attempted_queries,
),
)
def collect_location_candidates(
*,
name: str | None = None,
source: str | None = None,
source_id: str | None = None,
operator: str | None = None,
site: str | None = None,
city: str | None = None,
country: str | None = None,
organization: str | None = None,
record_id: int | None = None,
) -> tuple[list[LocationCandidate], list[str]]:
"""Run the full resolution chain and return ranked candidates with attempted queries.
The unused ``source`` / ``source_id`` / ``record_id`` arguments are kept
for backward compatibility with the API handler that calls this function.
"""
query = build_compute_center_location_query(
name=name,
source=source,
source_id=source_id,
operator=operator,
site=site,
city=city,
country=country,
organization=organization,
)
return COMPUTE_CENTER_COLLECTION_PIPELINE.collect_candidates(query)
def build_compute_center_location_query(
*,
name: str | None = None,
source: str | None = None,
source_id: str | None = None,
operator: str | None = None,
site: str | None = None,
city: str | None = None,
country: str | None = None,
organization: str | None = None,
) -> LocationQuery:
name_value = coerce_str(name)
context: dict[str, str] = {
"source": coerce_str(source),
"source_id": coerce_str(source_id),
"name": name_value,
"name_short": _short_system_name(name_value),
"city": coerce_str(city),
"country": coerce_str(country),
"site": coerce_str(site or organization),
"operator": coerce_str(operator or organization),
"organization": coerce_str(organization),
}
return _context_to_query(context)
def _record_operator(metadata: dict[str, Any]) -> str | None:
return coerce_str(
metadata.get("operator")
or metadata.get("organization")
or metadata.get("owner")
or metadata.get("manufacturer")
) or None
async def seed_compute_center_locations_from_source_coords(
session: AsyncSession,
) -> None:
"""Seed stored compute-center locations only from real source coordinates."""
stmt = (
select(CollectedData)
.where(CollectedData.source.in_(["top500", "epoch_ai_gpu"]))
.where(CollectedData.is_current.is_(True))
)
result = await session.execute(stmt)
records = result.scalars().all()
changed = False
for record in records:
source_value = coerce_str(getattr(record, "source", None))
source_id = coerce_str(getattr(record, "source_id", None))
if not source_value or not source_id:
continue
latitude = parse_float(get_record_field(record, "latitude"))
longitude = parse_float(get_record_field(record, "longitude"))
if latitude in (None, 0.0) or longitude in (None, 0.0):
continue
existing = await session.scalar(
select(ComputeCenterLocationRecord)
.where(ComputeCenterLocationRecord.source == source_value)
.where(ComputeCenterLocationRecord.source_id == source_id)
)
if existing:
continue
metadata = record.extra_data or {}
session.add(
ComputeCenterLocationRecord(
source=source_value,
source_id=source_id,
name=getattr(record, "name", None),
operator=_record_operator(metadata),
site=coerce_str(metadata.get("site") or metadata.get("organization")) or None,
city=coerce_str(get_record_field(record, "city")) or None,
country=coerce_str(get_record_field(record, "country")) or None,
latitude=latitude,
longitude=longitude,
precision="precise",
confidence=1.0,
location_source="source_coordinates",
source_note="Seeded from source-provided compute-center coordinates",
raw_payload={
"record_id": getattr(record, "id", None),
"source": source_value,
"source_id": source_id,
},
needs_confirmation=False,
verification_status="source_provided",
verified_at=None,
)
)
changed = True
if changed:
await session.commit()
await refresh_compute_center_location_cache(session)
async def upsert_compute_center_location(
session: AsyncSession,
*,
source: str,
source_id: str,
name: str | None = None,
operator: str | None = None,
site: str | None = None,
city: str | None = None,
country: str | None = None,
latitude: float,
longitude: float,
precision: str = "city",
confidence: float | None = None,
location_source: str = "manual_selection",
source_url: str | None = None,
source_note: str | None = None,
raw_payload: dict[str, Any] | None = None,
needs_confirmation: bool = False,
verification_status: str = "verified",
) -> ComputeCenterLocationRecord:
existing = await session.scalar(
select(ComputeCenterLocationRecord)
.where(ComputeCenterLocationRecord.source == source)
.where(ComputeCenterLocationRecord.source_id == source_id)
)
verified_at = None if needs_confirmation else datetime.now(UTC)
values = {
"name": name,
"operator": operator,
"site": site,
"city": city,
"country": country,
"latitude": latitude,
"longitude": longitude,
"precision": precision,
"confidence": confidence,
"location_source": location_source,
"source_url": source_url,
"source_note": source_note,
"raw_payload": raw_payload or {},
"needs_confirmation": needs_confirmation,
"verification_status": verification_status,
"verified_at": verified_at,
}
if existing:
for key, value in values.items():
setattr(existing, key, value)
record = existing
else:
record = ComputeCenterLocationRecord(
source=source,
source_id=source_id,
**values,
)
session.add(record)
await session.commit()
await session.refresh(record)
await refresh_compute_center_location_cache(session)
return record