887 lines
30 KiB
Python
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/plans/location-resolver-shared-pipeline-plan.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
|