Justin Fisher
1697 lines
63 KiB
Python
1697 lines
63 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Submit NSRDB polygon archive requests for county-level metrics.
|
|
|
|
This shared engine posts county polygon WKT to the NSRDB archive workflow and
|
|
records the download URLs returned by the API. Metric-specific wrappers supply
|
|
the requested attributes, artifact label, and output paths.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import csv
|
|
import getpass
|
|
import json
|
|
import re
|
|
import time
|
|
import urllib.error
|
|
import urllib.parse
|
|
import urllib.request
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from dataclasses import dataclass, field
|
|
from datetime import datetime, timedelta, timezone
|
|
from enum import Enum
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import geopandas as gpd
|
|
from shapely.geometry import box
|
|
from shapely.wkt import dumps as dump_wkt
|
|
|
|
|
|
DEFAULT_COUNTIES_GEOJSON = Path("data/geojson-counties-fips.json")
|
|
DEFAULT_ENDPOINT = "https://developer.nlr.gov/api/nsrdb/v2/solar/nsrdb-GOES-tmy-v4-0-0-download.json"
|
|
DEFAULT_POLAR_ENDPOINT = "https://developer.nlr.gov/api/nsrdb/v2/solar/nsrdb-polar-tmy-v4-0-0-download.json"
|
|
DEFAULT_SITE_COUNT_ENDPOINT = "https://developer.nlr.gov/api/nsrdb/v2/site-count.json"
|
|
DEFAULT_TMY_NAME = "tmy-2024"
|
|
DEFAULT_POLAR_TMY_NAME = "tmy"
|
|
DEFAULT_POLAR_MIN_LATITUDE = 60.0
|
|
DEFAULT_MAX_REQUEST_WEIGHT = 175000100
|
|
DEFAULT_SITE_COUNT_MAX_WKT_CHARS = 7000
|
|
DEFAULT_IN_FLIGHT_RETRIES = 3
|
|
DEFAULT_QUEUE_STATUS_TIMEOUT = 15
|
|
DEFAULT_QUEUE_STATUS_WORKERS = 8
|
|
DEFAULT_QUEUE_LOOKBACK_HOURS = 48
|
|
DEFAULT_RECENT_DOWNLOAD_WINDOW = 300.0
|
|
DEFAULT_SERVER_COUNT_GRACE = 1800.0
|
|
NSRDB_MAX_IN_FLIGHT_JOBS = 20
|
|
RECENT_DOWNLOAD_BYPASS_VISIBLE_SLOTS = 15
|
|
SMALL_POLYGON_WAIT = 60.0
|
|
MEDIUM_POLYGON_WAIT = 120.0
|
|
LARGE_POLYGON_WAIT = 360.0
|
|
EXCEPTIONAL_POLYGON_WAIT = 900.0
|
|
DEFAULT_TILE_MAX_WKT_CHARS = 6500
|
|
DEFAULT_MAX_TILES_PER_COUNTY = 64
|
|
DEFAULT_SKIP_INDEPENDENT_CITY_FIPS = "51530,51570,51678,51685,51830"
|
|
EMAIL_PATTERN = re.compile(r"[\w.!#$%&'*+/=?^`{|}~-]+@[\w-]+(?:\.[\w-]+)+")
|
|
SENSITIVE_QUERY_PATTERN = re.compile(r"((?:api_key|email)=)([^&\s,\"']+)", re.IGNORECASE)
|
|
|
|
STATE_FIPS_TO_ABBR = {
|
|
"01": "AL",
|
|
"02": "AK",
|
|
"04": "AZ",
|
|
"05": "AR",
|
|
"06": "CA",
|
|
"08": "CO",
|
|
"09": "CT",
|
|
"10": "DE",
|
|
"11": "DC",
|
|
"12": "FL",
|
|
"13": "GA",
|
|
"15": "HI",
|
|
"16": "ID",
|
|
"17": "IL",
|
|
"18": "IN",
|
|
"19": "IA",
|
|
"20": "KS",
|
|
"21": "KY",
|
|
"22": "LA",
|
|
"23": "ME",
|
|
"24": "MD",
|
|
"25": "MA",
|
|
"26": "MI",
|
|
"27": "MN",
|
|
"28": "MS",
|
|
"29": "MO",
|
|
"30": "MT",
|
|
"31": "NE",
|
|
"32": "NV",
|
|
"33": "NH",
|
|
"34": "NJ",
|
|
"35": "NM",
|
|
"36": "NY",
|
|
"37": "NC",
|
|
"38": "ND",
|
|
"39": "OH",
|
|
"40": "OK",
|
|
"41": "OR",
|
|
"42": "PA",
|
|
"44": "RI",
|
|
"45": "SC",
|
|
"46": "SD",
|
|
"47": "TN",
|
|
"48": "TX",
|
|
"49": "UT",
|
|
"50": "VT",
|
|
"51": "VA",
|
|
"53": "WA",
|
|
"54": "WV",
|
|
"55": "WI",
|
|
"56": "WY",
|
|
"60": "AS",
|
|
"66": "GU",
|
|
"69": "MP",
|
|
"72": "PR",
|
|
"78": "VI",
|
|
}
|
|
|
|
REQUEST_FIELDS = [
|
|
"county_fips",
|
|
"county_name",
|
|
"state_fips",
|
|
"state_abbr",
|
|
"profile_key",
|
|
"source",
|
|
"name",
|
|
"tile_id",
|
|
"tile_count",
|
|
"tile_area_fraction",
|
|
"site_count",
|
|
"request_weight",
|
|
"wkt_chars",
|
|
"status",
|
|
"message",
|
|
"download_url",
|
|
"response_json",
|
|
"submitted_at_utc",
|
|
]
|
|
|
|
ERROR_FIELDS = [
|
|
"county_fips",
|
|
"county_name",
|
|
"state_fips",
|
|
"state_abbr",
|
|
"profile_key",
|
|
"source",
|
|
"tile_id",
|
|
"tile_count",
|
|
"tile_area_fraction",
|
|
"site_count",
|
|
"request_weight",
|
|
"wkt_chars",
|
|
"error_type",
|
|
"status_code",
|
|
"request_url",
|
|
"response",
|
|
"message",
|
|
"logged_at_utc",
|
|
]
|
|
|
|
|
|
class CountyRequestState(Enum):
|
|
"""Distinct lifecycle phases for one county archive request."""
|
|
|
|
PENDING = "pending"
|
|
SUBMITTING = "submitting"
|
|
WAITING_TO_RETRY = "waiting_to_retry"
|
|
REQUESTED = "requested"
|
|
SKIPPED = "skipped"
|
|
FAILED = "failed"
|
|
|
|
|
|
class CountyRequestEvent(Enum):
|
|
"""Events that may move a county request to another lifecycle phase."""
|
|
|
|
START = "start"
|
|
EXISTING_REQUEST_FOUND = "existing_request_found"
|
|
REQUEST_SUCCEEDED = "request_succeeded"
|
|
RETRY_REQUIRED = "retry_required"
|
|
RETRY_STARTED = "retry_started"
|
|
COUNTY_SKIPPED = "county_skipped"
|
|
REQUEST_FAILED = "request_failed"
|
|
|
|
|
|
COUNTY_REQUEST_TRANSITIONS = {
|
|
(CountyRequestState.PENDING, CountyRequestEvent.START): CountyRequestState.SUBMITTING,
|
|
(
|
|
CountyRequestState.PENDING,
|
|
CountyRequestEvent.EXISTING_REQUEST_FOUND,
|
|
): CountyRequestState.SKIPPED,
|
|
(
|
|
CountyRequestState.SUBMITTING,
|
|
CountyRequestEvent.REQUEST_SUCCEEDED,
|
|
): CountyRequestState.REQUESTED,
|
|
(
|
|
CountyRequestState.SUBMITTING,
|
|
CountyRequestEvent.RETRY_REQUIRED,
|
|
): CountyRequestState.WAITING_TO_RETRY,
|
|
(
|
|
CountyRequestState.WAITING_TO_RETRY,
|
|
CountyRequestEvent.RETRY_STARTED,
|
|
): CountyRequestState.SUBMITTING,
|
|
(
|
|
CountyRequestState.SUBMITTING,
|
|
CountyRequestEvent.COUNTY_SKIPPED,
|
|
): CountyRequestState.SKIPPED,
|
|
(
|
|
CountyRequestState.SUBMITTING,
|
|
CountyRequestEvent.REQUEST_FAILED,
|
|
): CountyRequestState.FAILED,
|
|
}
|
|
|
|
|
|
class InvalidCountyRequestTransition(RuntimeError):
|
|
"""Reject an event that is not valid for the current request state."""
|
|
|
|
|
|
@dataclass
|
|
class CountyRequestStateMachine:
|
|
"""Track and validate the lifecycle of one county archive request."""
|
|
|
|
label: str
|
|
state: CountyRequestState = CountyRequestState.PENDING
|
|
history: list[CountyRequestState] = field(
|
|
default_factory=lambda: [CountyRequestState.PENDING]
|
|
)
|
|
|
|
def transition(self, event: CountyRequestEvent) -> CountyRequestState:
|
|
"""Apply one event and return the resulting state."""
|
|
next_state = COUNTY_REQUEST_TRANSITIONS.get((self.state, event))
|
|
if next_state is None:
|
|
raise InvalidCountyRequestTransition(
|
|
f"Invalid county request transition for {self.label}: "
|
|
f"{event.value} while {self.state.value}."
|
|
)
|
|
self.state = next_state
|
|
self.history.append(next_state)
|
|
return next_state
|
|
|
|
|
|
def normalize_fips(value: object, width: int) -> str:
|
|
"""Return a zero-padded FIPS string from a mixed text or numeric value."""
|
|
digits = "".join(character for character in str(value).strip() if character.isdigit())
|
|
return digits.zfill(width)[-width:] if digits else ""
|
|
|
|
|
|
def prompt_for_secret(prompt: str, current_value: str | None) -> str:
|
|
"""Prompt for a sensitive value only when it was not supplied."""
|
|
if current_value:
|
|
return current_value
|
|
return getpass.getpass(prompt).strip()
|
|
|
|
|
|
def prompt_for_text(prompt: str, current_value: str | None) -> str:
|
|
"""Prompt for normal text only when it was not supplied."""
|
|
if current_value:
|
|
return current_value
|
|
return input(prompt).strip()
|
|
|
|
|
|
def redact_url(url: str) -> str:
|
|
"""Return a request URL with sensitive query values removed."""
|
|
parsed_url = urllib.parse.urlsplit(url)
|
|
query = urllib.parse.parse_qsl(parsed_url.query, keep_blank_values=True)
|
|
redacted_query = [
|
|
(key, "<redacted>" if key in {"api_key", "email"} else value)
|
|
for key, value in query
|
|
]
|
|
return urllib.parse.urlunsplit(
|
|
parsed_url._replace(query=urllib.parse.urlencode(redacted_query))
|
|
)
|
|
|
|
|
|
def redact_sensitive_text(text: str) -> str:
|
|
"""Remove likely credentials and email addresses from log text."""
|
|
without_query_values = SENSITIVE_QUERY_PATTERN.sub(r"\1<redacted>", text)
|
|
return EMAIL_PATTERN.sub("<redacted>", without_query_values)
|
|
|
|
|
|
def json_error_messages(body: str) -> list[str]:
|
|
"""Extract API error messages from a JSON response body when possible."""
|
|
try:
|
|
payload = json.loads(body)
|
|
except json.JSONDecodeError:
|
|
return []
|
|
|
|
errors = payload.get("errors") if isinstance(payload, dict) else None
|
|
if isinstance(errors, list):
|
|
return [redact_sensitive_text(str(error)) for error in errors]
|
|
if isinstance(errors, str):
|
|
return [redact_sensitive_text(errors)]
|
|
|
|
error = payload.get("error") if isinstance(payload, dict) else None
|
|
if isinstance(error, dict):
|
|
message = error.get("message") or error.get("code")
|
|
return [redact_sensitive_text(str(message))] if message else []
|
|
|
|
return []
|
|
|
|
|
|
def http_error_message(error: urllib.error.HTTPError, url: str, body: str) -> str:
|
|
"""Build a sanitized HTTP error message with parsed API errors first."""
|
|
body_excerpt = redact_sensitive_text(body[:1000]) if body else "No response body returned."
|
|
retry_after = error.headers.get("Retry-After") if error.headers else None
|
|
retry_message = f"\n Retry-After: {retry_after}" if retry_after else ""
|
|
api_errors = json_error_messages(body)
|
|
api_error_message = f"\n API errors: {' | '.join(api_errors)}" if api_errors else ""
|
|
return f"HTTP {error.code} {error.reason}\n URL: {redact_url(url)}{retry_message}{api_error_message}\n Response: {body_excerpt}"
|
|
|
|
|
|
class NsrdArchiveRequestError(RuntimeError):
|
|
"""Store a failed NSRDB archive request with sanitized context."""
|
|
|
|
def __init__(
|
|
self,
|
|
message: str,
|
|
status_code: int | None = None,
|
|
url: str | None = None,
|
|
response: str = "",
|
|
retry_after: str | None = None,
|
|
) -> None:
|
|
super().__init__(message)
|
|
self.status_code = status_code
|
|
self.url = url
|
|
self.response = response
|
|
self.retry_after = retry_after
|
|
|
|
|
|
class IndependentCitySkip(RuntimeError):
|
|
"""Mark an intentionally skipped independent city with no polygon grid sites."""
|
|
|
|
|
|
class LocalQueueCapacityError(RuntimeError):
|
|
"""Stop before submission when locally tracked NSRDB jobs fill the queue."""
|
|
|
|
|
|
class RequestRateLimiter:
|
|
"""Coordinate request starts so archive jobs respect NSRDB limits."""
|
|
|
|
def __init__(self, delay: float) -> None:
|
|
self.delay = max(0.0, delay)
|
|
self._next_request_at = 0.0
|
|
|
|
def wait(self) -> None:
|
|
"""Wait until the next request can start."""
|
|
if self.delay <= 0:
|
|
return
|
|
|
|
now = time.monotonic()
|
|
wait_seconds = max(0.0, self._next_request_at - now)
|
|
self._next_request_at = max(now, self._next_request_at) + self.delay
|
|
if wait_seconds:
|
|
time.sleep(wait_seconds)
|
|
|
|
|
|
class ArchiveQueueMonitor:
|
|
"""Track recent NSRDB jobs without submitting another archive request."""
|
|
|
|
def __init__(self, manifest_dir: Path, timeout: int, workers: int) -> None:
|
|
self.manifest_dir = manifest_dir
|
|
self.timeout = timeout
|
|
self.workers = workers
|
|
self.ready_urls: set[str] = set()
|
|
self.pending_urls: set[str] = set()
|
|
self.recently_ready_at: dict[str, float] = {}
|
|
self.paced_submissions = 0
|
|
self.filling_open_queue = False
|
|
self.unobserved_jobs = 0
|
|
self.unobserved_jobs_until = 0.0
|
|
self.jobs, self.tracked_site_counts = self._load_recent_jobs()
|
|
|
|
@staticmethod
|
|
def _downloaded_archive_path(response_path: Path) -> Path:
|
|
"""Return the expected downloaded ZIP path for a response JSON path."""
|
|
archive_dir_name = response_path.parent.name.replace(
|
|
"_request_responses",
|
|
"_archives",
|
|
)
|
|
archive_name = response_path.stem.removesuffix("_response") + ".zip"
|
|
return response_path.parent.parent / archive_dir_name / archive_name
|
|
|
|
def _response_path_from_row(self, row: dict[str, str]) -> Path | None:
|
|
"""Resolve a manifest response path when one was recorded."""
|
|
response_json = row.get("response_json", "").strip()
|
|
if not response_json:
|
|
return None
|
|
|
|
response_path = Path(response_json)
|
|
if response_path.is_absolute():
|
|
return response_path
|
|
|
|
if response_path.exists():
|
|
return response_path
|
|
|
|
manifest_relative = self.manifest_dir / response_path
|
|
if manifest_relative.exists():
|
|
return manifest_relative
|
|
|
|
return response_path
|
|
|
|
def _archive_was_downloaded(self, response_path: Path | None) -> bool:
|
|
"""Return whether local output proves an archive job completed."""
|
|
if response_path is None:
|
|
return False
|
|
return self._downloaded_archive_path(response_path).is_file()
|
|
|
|
def _load_recent_jobs(self) -> tuple[dict[str, str], dict[str, int]]:
|
|
"""Load plausible in-flight jobs and their polygon sizes."""
|
|
cutoff = datetime.now(timezone.utc) - timedelta(hours=DEFAULT_QUEUE_LOOKBACK_HOURS)
|
|
jobs: dict[str, str] = {}
|
|
site_counts: dict[str, int] = {}
|
|
for manifest_path in sorted(self.manifest_dir.glob("*request_manifest.csv")):
|
|
for row in read_existing_rows(manifest_path).values():
|
|
url = row.get("download_url", "").strip()
|
|
submitted_at = row.get("submitted_at_utc", "").strip()
|
|
if not url or not submitted_at:
|
|
continue
|
|
try:
|
|
submitted = datetime.fromisoformat(submitted_at.replace("Z", "+00:00"))
|
|
except ValueError:
|
|
continue
|
|
if submitted.tzinfo is None:
|
|
submitted = submitted.replace(tzinfo=timezone.utc)
|
|
if submitted < cutoff:
|
|
continue
|
|
if self._archive_was_downloaded(self._response_path_from_row(row)):
|
|
continue
|
|
tile_id = row.get("tile_id", "").strip()
|
|
label = row.get("county_fips", "").strip()
|
|
jobs[url] = f"{label}:{tile_id}" if tile_id else label
|
|
try:
|
|
site_counts[url] = int(row.get("site_count", ""))
|
|
except (TypeError, ValueError):
|
|
pass
|
|
|
|
for response_path in sorted(self.manifest_dir.glob("polygon*_request_responses/*.json")):
|
|
modified = datetime.fromtimestamp(response_path.stat().st_mtime, timezone.utc)
|
|
if modified < cutoff:
|
|
continue
|
|
if self._archive_was_downloaded(response_path):
|
|
continue
|
|
try:
|
|
response = json.loads(response_path.read_text(encoding="utf-8"))
|
|
except (OSError, json.JSONDecodeError):
|
|
continue
|
|
outputs = response.get("outputs") if isinstance(response, dict) else None
|
|
url = outputs.get("downloadUrl", "").strip() if isinstance(outputs, dict) else ""
|
|
if url:
|
|
jobs.setdefault(url, response_path.stem)
|
|
return jobs, site_counts
|
|
|
|
def record_submission(self, row: dict[str, str]) -> None:
|
|
"""Immediately count a newly accepted county or tile as in flight."""
|
|
url = row.get("download_url", "").strip()
|
|
if not url:
|
|
return
|
|
tile_id = row.get("tile_id", "").strip()
|
|
label = row.get("county_fips", "").strip()
|
|
self.jobs[url] = f"{label}:{tile_id}" if tile_id else label
|
|
self.pending_urls.add(url)
|
|
try:
|
|
self.tracked_site_counts[url] = int(row.get("site_count", ""))
|
|
except (TypeError, ValueError):
|
|
pass
|
|
|
|
def dynamic_wait_seconds(self) -> float:
|
|
"""Choose queue pacing from the largest active or recently ready polygon.
|
|
|
|
Large county polygons and Alaska tiles can take much longer for NLR to
|
|
prepare. Sizes remain tracked through the five-minute recent-download
|
|
window so the next status check reflects the jobs still shaping queue
|
|
throughput, without exposing the tracker in console output.
|
|
"""
|
|
largest_site_count = max(self.tracked_site_counts.values(), default=0)
|
|
if largest_site_count >= 551:
|
|
return EXCEPTIONAL_POLYGON_WAIT
|
|
if largest_site_count >= 251:
|
|
return LARGE_POLYGON_WAIT
|
|
if largest_site_count >= 111:
|
|
return MEDIUM_POLYGON_WAIT
|
|
return SMALL_POLYGON_WAIT
|
|
|
|
def wait_for_dynamic_queue(self, cooldown: float, resume_message: str) -> bool:
|
|
"""Wait for queue progress, checking long waits for newly open capacity."""
|
|
if cooldown == LARGE_POLYGON_WAIT:
|
|
wait_segments = (180.0, 180.0)
|
|
elif cooldown == EXCEPTIONAL_POLYGON_WAIT:
|
|
wait_segments = (300.0, 300.0, 300.0)
|
|
else:
|
|
wait_for_queue(cooldown, resume_message)
|
|
return False
|
|
|
|
for segment_index, wait_seconds in enumerate(wait_segments):
|
|
if wait_seconds > 0:
|
|
time.sleep(wait_seconds)
|
|
if segment_index == len(wait_segments) - 1:
|
|
break
|
|
|
|
counts = self.summarize()
|
|
occupied = (
|
|
counts["pending"]
|
|
+ counts["unobserved"]
|
|
+ counts["unknown"]
|
|
)
|
|
visible_slots = max(0, NSRDB_MAX_IN_FLIGHT_JOBS - occupied)
|
|
if visible_slots >= RECENT_DOWNLOAD_BYPASS_VISIBLE_SLOTS:
|
|
self.filling_open_queue = True
|
|
self.paced_submissions = 0
|
|
print(
|
|
f"Queue checkpoint shows {visible_slots} slots visible; "
|
|
"ending the wait and entering fill mode."
|
|
)
|
|
return True
|
|
print(
|
|
f"Queue checkpoint shows {visible_slots} slots visible; "
|
|
"continuing the dynamic wait."
|
|
)
|
|
|
|
print(resume_message)
|
|
return False
|
|
|
|
def observe_server_count(self, message: str) -> None:
|
|
"""Infer jobs absent from local URL tracking from an NLR response."""
|
|
match = re.search(r"owns\s+(\d+)\s+in-flight jobs", message, re.IGNORECASE)
|
|
if not match:
|
|
return
|
|
|
|
server_count = int(match.group(1))
|
|
counts = self.summarize()
|
|
locally_observed = counts["pending"] + counts["unknown"]
|
|
unobserved = max(0, server_count - locally_observed)
|
|
if unobserved:
|
|
self.unobserved_jobs = max(self.unobserved_jobs, unobserved)
|
|
self.unobserved_jobs_until = (
|
|
time.monotonic() + DEFAULT_SERVER_COUNT_GRACE
|
|
)
|
|
|
|
def summarize(self) -> dict[str, int]:
|
|
"""Probe existing archive URLs and return queue-status counts."""
|
|
urls = set(self.jobs)
|
|
statuses: list[str] = []
|
|
ambiguous_403 = 0
|
|
unresolved_urls = sorted(urls - self.ready_urls)
|
|
if unresolved_urls:
|
|
with ThreadPoolExecutor(max_workers=self.workers) as executor:
|
|
statuses = list(
|
|
executor.map(
|
|
lambda url: archive_download_status(url, self.timeout),
|
|
unresolved_urls,
|
|
)
|
|
)
|
|
now = time.monotonic()
|
|
for url, status in zip(unresolved_urls, statuses):
|
|
if status == "pending":
|
|
if url not in self.pending_urls:
|
|
ambiguous_403 += 1
|
|
continue
|
|
if status != "ready":
|
|
continue
|
|
if url in self.pending_urls:
|
|
self.recently_ready_at[url] = now
|
|
else:
|
|
self.tracked_site_counts.pop(url, None)
|
|
self.pending_urls.discard(url)
|
|
self.ready_urls.add(url)
|
|
|
|
now = time.monotonic()
|
|
expired_recent_urls = {
|
|
url
|
|
for url, ready_at in self.recently_ready_at.items()
|
|
if now - ready_at > DEFAULT_RECENT_DOWNLOAD_WINDOW
|
|
}
|
|
self.recently_ready_at = {
|
|
url: ready_at
|
|
for url, ready_at in self.recently_ready_at.items()
|
|
if now - ready_at <= DEFAULT_RECENT_DOWNLOAD_WINDOW
|
|
}
|
|
for url in expired_recent_urls:
|
|
self.tracked_site_counts.pop(url, None)
|
|
recently_downloaded = len(urls & self.recently_ready_at.keys())
|
|
if self.unobserved_jobs_until <= now:
|
|
self.unobserved_jobs = 0
|
|
|
|
return {
|
|
"pending": sum(
|
|
status == "pending" and url in self.pending_urls
|
|
for url, status in zip(unresolved_urls, statuses)
|
|
),
|
|
"ambiguous_403": ambiguous_403,
|
|
"recently_downloaded": recently_downloaded,
|
|
"unobserved": self.unobserved_jobs,
|
|
"ready": len(urls & self.ready_urls),
|
|
"unknown": statuses.count("unknown"),
|
|
"total": len(urls),
|
|
}
|
|
|
|
def wait_for_capacity(
|
|
self,
|
|
next_label: str,
|
|
retries: int,
|
|
) -> None:
|
|
"""Wait until a slot is visible before allowing another archive POST."""
|
|
queue_attempt = 0
|
|
while True:
|
|
counts = self.summarize()
|
|
occupied = (
|
|
counts["pending"]
|
|
+ counts["unobserved"]
|
|
+ counts["unknown"]
|
|
)
|
|
visible_slots = max(0, NSRDB_MAX_IN_FLIGHT_JOBS - occupied)
|
|
recently_downloaded = counts["recently_downloaded"]
|
|
paced_limit = max(1, recently_downloaded // 2) if recently_downloaded else 0
|
|
if visible_slots >= RECENT_DOWNLOAD_BYPASS_VISIBLE_SLOTS:
|
|
# A substantially cleared queue should be refilled before recent
|
|
# completions are allowed to trigger another pacing delay.
|
|
self.filling_open_queue = True
|
|
self.paced_submissions = 0
|
|
if not recently_downloaded:
|
|
self.paced_submissions = 0
|
|
print(
|
|
f"NSRDB queue before {next_label}: "
|
|
f"{counts['pending']} pending, "
|
|
f"{recently_downloaded} downloaded in the last 5 minutes, "
|
|
f"{counts['unobserved']} unobserved (inferred from NLR's total), "
|
|
f"{counts.get('ambiguous_403', 0)} historical S3 403, "
|
|
f"{counts['unknown']} unknown, "
|
|
f"{visible_slots} slots visible."
|
|
)
|
|
|
|
if occupied >= NSRDB_MAX_IN_FLIGHT_JOBS:
|
|
self.filling_open_queue = False
|
|
if queue_attempt >= retries:
|
|
break
|
|
queue_attempt += 1
|
|
cooldown = self.dynamic_wait_seconds()
|
|
print(
|
|
f"Queue monitor shows {occupied} possible in-flight county/tile jobs; "
|
|
f"waiting {cooldown:.0f} seconds before checking existing download URLs again."
|
|
)
|
|
if self.wait_for_dynamic_queue(
|
|
cooldown,
|
|
"Queue wait complete; checking capacity again.",
|
|
):
|
|
return
|
|
continue
|
|
|
|
if self.filling_open_queue:
|
|
return
|
|
|
|
if not paced_limit or self.paced_submissions < paced_limit:
|
|
if paced_limit:
|
|
self.paced_submissions += 1
|
|
return
|
|
|
|
cooldown = self.dynamic_wait_seconds()
|
|
print(
|
|
f"Submitted {self.paced_submissions} jobs after {recently_downloaded} recent "
|
|
f"downloads; waiting {cooldown:.0f} seconds "
|
|
"before starting the next paced batch."
|
|
)
|
|
if self.wait_for_dynamic_queue(
|
|
cooldown,
|
|
"Pacing wait complete; refreshing archive status.",
|
|
):
|
|
return
|
|
self.paced_submissions = 0
|
|
queue_attempt = 0
|
|
|
|
raise LocalQueueCapacityError(
|
|
f"Stopped before submitting {next_label}: the queue monitor still shows "
|
|
f"{occupied} possible in-flight jobs after {retries} retries."
|
|
)
|
|
|
|
|
|
def load_counties(counties_geojson: Path, include_puerto_rico: bool) -> list[dict[str, Any]]:
|
|
"""Load county polygons and metadata in EPSG:4326."""
|
|
counties = gpd.read_file(counties_geojson)
|
|
if counties.crs is None:
|
|
counties = counties.set_crs("EPSG:4326")
|
|
else:
|
|
counties = counties.to_crs("EPSG:4326")
|
|
|
|
rows: list[dict[str, Any]] = []
|
|
for _, county in counties.iterrows():
|
|
county_fips = (
|
|
normalize_fips(county.get("id"), 5)
|
|
or normalize_fips(county.get("GEOID"), 5)
|
|
or normalize_fips(county.get("GEOID10"), 5)
|
|
)
|
|
if not county_fips:
|
|
state_fips = normalize_fips(county.get("STATE"), 2)
|
|
county_code = normalize_fips(county.get("COUNTY"), 3)
|
|
county_fips = normalize_fips(f"{state_fips}{county_code}", 5)
|
|
|
|
if not county_fips:
|
|
continue
|
|
|
|
state_fips = county_fips[:2]
|
|
if state_fips == "72" and not include_puerto_rico:
|
|
continue
|
|
|
|
geometry = county.geometry
|
|
if geometry is None or geometry.is_empty:
|
|
continue
|
|
if not geometry.is_valid:
|
|
geometry = geometry.buffer(0)
|
|
|
|
point = geometry.representative_point()
|
|
rows.append(
|
|
{
|
|
"county_fips": county_fips,
|
|
"county_name": str(county.get("NAME") or county.get("name") or f"County {county_fips}").strip(),
|
|
"state_fips": state_fips,
|
|
"state_abbr": STATE_FIPS_TO_ABBR.get(state_fips, state_fips),
|
|
"lat": float(point.y),
|
|
"lon": float(point.x),
|
|
"geometry": geometry,
|
|
}
|
|
)
|
|
|
|
return sorted(rows, key=lambda row: row["county_fips"])
|
|
|
|
|
|
def slice_batch(rows: list[dict[str, Any]], start: int, limit: int | None) -> list[dict[str, Any]]:
|
|
"""Return the requested county batch."""
|
|
if start < 0:
|
|
raise ValueError("--start must be 0 or greater.")
|
|
if limit is None:
|
|
return rows[start:]
|
|
if limit < 1:
|
|
raise ValueError("--limit must be 1 or greater.")
|
|
return rows[start : start + limit]
|
|
|
|
|
|
def county_wkt(county: dict[str, Any], precision: int) -> str:
|
|
"""Convert a county geometry to compact WKT for the NSRDB API."""
|
|
return geometry_wkt(county["geometry"], precision)
|
|
|
|
|
|
def geometry_wkt(geometry: Any, precision: int) -> str:
|
|
"""Convert a geometry to compact WKT for the NSRDB API."""
|
|
return dump_wkt(geometry, rounding_precision=precision, trim=True)
|
|
|
|
|
|
def csv_row_key(row: dict[str, str]) -> str:
|
|
"""Return the manifest/error key for a whole-county or tile row."""
|
|
county_fips = row.get("county_fips", "")
|
|
tile_id = row.get("tile_id") or "county"
|
|
return f"{county_fips}:{tile_id}"
|
|
|
|
|
|
def read_existing_rows(csv_path: Path) -> dict[str, dict[str, str]]:
|
|
"""Read an existing county/tile-keyed CSV manifest."""
|
|
if not csv_path.exists():
|
|
return {}
|
|
with csv_path.open(newline="", encoding="utf-8") as handle:
|
|
return {
|
|
csv_row_key(row): row
|
|
for row in csv.DictReader(handle)
|
|
if row.get("county_fips")
|
|
}
|
|
|
|
|
|
def write_rows(rows_by_fips: dict[str, dict[str, str]], csv_path: Path, fields: list[str]) -> None:
|
|
"""Write a county-keyed CSV manifest."""
|
|
csv_path.parent.mkdir(parents=True, exist_ok=True)
|
|
with csv_path.open("w", newline="", encoding="utf-8") as handle:
|
|
writer = csv.DictWriter(handle, fieldnames=fields, extrasaction="ignore")
|
|
writer.writeheader()
|
|
for row_key in sorted(rows_by_fips):
|
|
writer.writerow(rows_by_fips[row_key])
|
|
|
|
|
|
def county_row_keys(rows_by_key: dict[str, dict[str, str]], county_fips: str) -> list[str]:
|
|
"""Return all existing manifest/error keys for a county."""
|
|
return [
|
|
row_key
|
|
for row_key, row in rows_by_key.items()
|
|
if row.get("county_fips") == county_fips
|
|
]
|
|
|
|
|
|
def build_url(endpoint: str, api_key: str) -> str:
|
|
"""Add the required api_key query parameter to an NSRDB endpoint."""
|
|
separator = "&" if "?" in endpoint else "?"
|
|
return f"{endpoint}{separator}{urllib.parse.urlencode({'api_key': api_key})}"
|
|
|
|
|
|
def post_form_json(
|
|
endpoint: str,
|
|
api_key: str,
|
|
payload: dict[str, str],
|
|
timeout: int,
|
|
rate_limiter: RequestRateLimiter,
|
|
) -> dict[str, Any]:
|
|
"""POST form data and parse a JSON response."""
|
|
url = build_url(endpoint, api_key)
|
|
data = urllib.parse.urlencode(payload).encode("utf-8")
|
|
request = urllib.request.Request(
|
|
url,
|
|
data=data,
|
|
headers={
|
|
"User-Agent": "county-climate-explorer/0.1",
|
|
"Content-Type": "application/x-www-form-urlencoded",
|
|
"Cache-Control": "no-cache",
|
|
},
|
|
method="POST",
|
|
)
|
|
try:
|
|
rate_limiter.wait()
|
|
with urllib.request.urlopen(request, timeout=timeout) as response:
|
|
body = response.read().decode("utf-8-sig")
|
|
except urllib.error.HTTPError as error:
|
|
body = error.read().decode("utf-8-sig", errors="replace").strip()
|
|
body_excerpt = redact_sensitive_text(body[:1000]) if body else "No response body returned."
|
|
retry_after = error.headers.get("Retry-After") if error.headers else None
|
|
raise NsrdArchiveRequestError(
|
|
http_error_message(error, url, body),
|
|
error.code,
|
|
redact_url(url),
|
|
body_excerpt,
|
|
retry_after,
|
|
) from error
|
|
|
|
try:
|
|
return json.loads(body)
|
|
except json.JSONDecodeError as exc:
|
|
raise NsrdArchiveRequestError(
|
|
f"Expected JSON response from {redact_url(url)}, got: {redact_sensitive_text(body[:1000])}",
|
|
url=redact_url(url),
|
|
response=redact_sensitive_text(body[:1000]),
|
|
) from exc
|
|
|
|
|
|
def get_query_json(
|
|
endpoint: str,
|
|
api_key: str,
|
|
query: dict[str, str],
|
|
timeout: int,
|
|
rate_limiter: RequestRateLimiter,
|
|
) -> dict[str, Any]:
|
|
"""GET query-string parameters and parse a JSON response."""
|
|
parsed_url = urllib.parse.urlsplit(endpoint)
|
|
existing_query = urllib.parse.parse_qsl(parsed_url.query, keep_blank_values=True)
|
|
full_query = existing_query + [("api_key", api_key)] + list(query.items())
|
|
url = urllib.parse.urlunsplit(
|
|
parsed_url._replace(query=urllib.parse.urlencode(full_query))
|
|
)
|
|
request = urllib.request.Request(
|
|
url,
|
|
headers={"User-Agent": "county-climate-explorer/0.1"},
|
|
method="GET",
|
|
)
|
|
try:
|
|
rate_limiter.wait()
|
|
with urllib.request.urlopen(request, timeout=timeout) as response:
|
|
body = response.read().decode("utf-8-sig")
|
|
except urllib.error.HTTPError as error:
|
|
body = error.read().decode("utf-8-sig", errors="replace").strip()
|
|
body_excerpt = redact_sensitive_text(body[:1000]) if body else "No response body returned."
|
|
retry_after = error.headers.get("Retry-After") if error.headers else None
|
|
raise NsrdArchiveRequestError(
|
|
http_error_message(error, url, body),
|
|
error.code,
|
|
redact_url(url),
|
|
body_excerpt,
|
|
retry_after,
|
|
) from error
|
|
|
|
try:
|
|
return json.loads(body)
|
|
except json.JSONDecodeError as exc:
|
|
raise NsrdArchiveRequestError(
|
|
f"Expected JSON response from {redact_url(url)}, got: {redact_sensitive_text(body[:1000])}",
|
|
url=redact_url(url),
|
|
response=redact_sensitive_text(body[:1000]),
|
|
) from exc
|
|
|
|
|
|
def request_profiles(args: argparse.Namespace, county: dict[str, Any]) -> list[dict[str, str]]:
|
|
"""Build the ordered list of NSRDB datasets to try for a county."""
|
|
goes = {
|
|
"profile_key": "goes-tmy",
|
|
"source": "NSRDB GOES TMY PSM v4 county polygon",
|
|
"dataset_key": "nsrdb-GOES-tmy-v4-0-0",
|
|
"endpoint": args.endpoint,
|
|
"name": args.name,
|
|
}
|
|
polar = {
|
|
"profile_key": "polar-tmy",
|
|
"source": "NSRDB Polar TMY PSM v4 county polygon",
|
|
"dataset_key": "nsrdb-polar-tmy-v4-0-0",
|
|
"endpoint": args.polar_endpoint,
|
|
"name": args.polar_name,
|
|
}
|
|
|
|
if args.profile == "goes":
|
|
return [goes]
|
|
if args.profile == "polar":
|
|
return [polar]
|
|
if args.profile == "auto" and args.polar_fallback and county["lat"] >= args.polar_min_latitude:
|
|
return [goes, polar]
|
|
return [goes]
|
|
|
|
|
|
def get_site_count(
|
|
args: argparse.Namespace,
|
|
api_key: str,
|
|
profile: dict[str, str],
|
|
wkt: str,
|
|
rate_limiter: RequestRateLimiter,
|
|
) -> int | None:
|
|
"""Return NSRDB site count for one profile, or None when disabled."""
|
|
if args.skip_site_count:
|
|
return None
|
|
if args.site_count_max_wkt_chars and len(wkt) > args.site_count_max_wkt_chars:
|
|
return None
|
|
response = get_query_json(
|
|
endpoint=args.site_count_endpoint,
|
|
api_key=api_key,
|
|
query={"wkt": wkt},
|
|
timeout=args.timeout,
|
|
rate_limiter=rate_limiter,
|
|
)
|
|
errors = response.get("errors") or []
|
|
if errors:
|
|
raise NsrdArchiveRequestError(f"Site-count API returned errors: {errors}")
|
|
outputs = response.get("outputs") or {}
|
|
value = outputs.get(profile["dataset_key"])
|
|
return int(value) if value is not None else 0
|
|
|
|
|
|
def request_weight(site_count: int | None, attributes: str, names: str, interval: int) -> int | None:
|
|
"""Calculate NSRDB request weight when site count is available."""
|
|
if site_count is None:
|
|
return None
|
|
attribute_count = len([value for value in attributes.split(",") if value.strip()])
|
|
year_count = len([value for value in names.split(",") if value.strip()])
|
|
intervals_per_year = int((60 / interval) * 24 * 365)
|
|
return site_count * attribute_count * year_count * intervals_per_year
|
|
|
|
|
|
def parse_fips_set(value: str | None) -> set[str]:
|
|
"""Parse a comma-delimited FIPS list."""
|
|
if not value:
|
|
return set()
|
|
return {
|
|
normalize_fips(item, 5)
|
|
for item in value.split(",")
|
|
if normalize_fips(item, 5)
|
|
}
|
|
|
|
|
|
def should_skip_independent_city(args: argparse.Namespace, county: dict[str, Any]) -> bool:
|
|
"""Return whether this county is an intentionally skipped independent city."""
|
|
return county["county_fips"] in parse_fips_set(args.skip_independent_city_fips)
|
|
|
|
|
|
def is_request_size_error(error: Exception) -> bool:
|
|
"""Return whether an API error means the county should be tiled."""
|
|
return isinstance(error, NsrdArchiveRequestError) and "request size exceeds maximum limit" in str(error).lower()
|
|
|
|
|
|
def clean_geometry(geometry: Any) -> Any:
|
|
"""Return a valid non-empty geometry when possible."""
|
|
if geometry is None or geometry.is_empty:
|
|
return geometry
|
|
if not geometry.is_valid:
|
|
geometry = geometry.buffer(0)
|
|
return geometry
|
|
|
|
|
|
def split_geometry(geometry: Any) -> list[Any]:
|
|
"""Split a geometry along the longer side of its bounds."""
|
|
min_x, min_y, max_x, max_y = geometry.bounds
|
|
if max_x <= min_x or max_y <= min_y:
|
|
return []
|
|
|
|
if (max_x - min_x) >= (max_y - min_y):
|
|
mid_x = (min_x + max_x) / 2
|
|
cells = [
|
|
box(min_x, min_y, mid_x, max_y),
|
|
box(mid_x, min_y, max_x, max_y),
|
|
]
|
|
else:
|
|
mid_y = (min_y + max_y) / 2
|
|
cells = [
|
|
box(min_x, min_y, max_x, mid_y),
|
|
box(min_x, mid_y, max_x, max_y),
|
|
]
|
|
|
|
parts: list[Any] = []
|
|
for cell in cells:
|
|
part = clean_geometry(geometry.intersection(cell))
|
|
if part is not None and not part.is_empty:
|
|
parts.append(part)
|
|
return parts
|
|
|
|
|
|
def geometry_area_fraction(geometry: Any, parent_geometry: Any) -> float:
|
|
"""Return an approximate geometry area fraction for logging."""
|
|
parent_area = float(parent_geometry.area or 0)
|
|
if parent_area <= 0:
|
|
return 0.0
|
|
return float(geometry.area) / parent_area
|
|
|
|
|
|
def archive_payload(args: argparse.Namespace, email: str, profile: dict[str, str], wkt: str) -> dict[str, str]:
|
|
"""Build the NSRDB archive request payload."""
|
|
payload = {
|
|
"attributes": args.attributes,
|
|
"names": profile["name"],
|
|
"utc": str(args.utc).lower(),
|
|
"leap_day": str(args.leap_day).lower(),
|
|
"interval": str(args.interval),
|
|
"email": email,
|
|
"wkt": wkt,
|
|
}
|
|
optional_fields = {
|
|
"full_name": args.full_name,
|
|
"affiliation": args.affiliation,
|
|
"reason": args.reason,
|
|
"mailing_list": str(args.mailing_list).lower() if args.mailing_list else None,
|
|
}
|
|
payload.update({key: value for key, value in optional_fields.items() if value})
|
|
return payload
|
|
|
|
|
|
def response_path(
|
|
args: argparse.Namespace,
|
|
county_fips: str,
|
|
profile_key: str,
|
|
name: str,
|
|
tile_id: str = "",
|
|
) -> Path:
|
|
"""Return the cached API response path for one request."""
|
|
tile_part = f"_{tile_id}" if tile_id else ""
|
|
path = (
|
|
args.response_dir
|
|
/ f"{county_fips}{tile_part}_{profile_key}_{name}_{args.artifact_label}_response.json"
|
|
)
|
|
if args.legacy_artifact_label and not path.exists():
|
|
legacy_path = (
|
|
args.response_dir
|
|
/ f"{county_fips}{tile_part}_{profile_key}_{name}_{args.legacy_artifact_label}_response.json"
|
|
)
|
|
if legacy_path.exists():
|
|
return legacy_path
|
|
return path
|
|
|
|
|
|
def build_success_row(
|
|
county: dict[str, Any],
|
|
profile: dict[str, str],
|
|
site_count: int | None,
|
|
weight: int | None,
|
|
wkt: str,
|
|
response: dict[str, Any],
|
|
response_json: Path,
|
|
tile_id: str = "",
|
|
tile_count: int | None = None,
|
|
tile_area_fraction: float | None = None,
|
|
) -> dict[str, str]:
|
|
"""Build a request manifest row from an archive acknowledgement."""
|
|
outputs = response.get("outputs") or {}
|
|
return {
|
|
"county_fips": county["county_fips"],
|
|
"county_name": county["county_name"],
|
|
"state_fips": county["state_fips"],
|
|
"state_abbr": county["state_abbr"],
|
|
"profile_key": profile["profile_key"],
|
|
"source": profile["source"],
|
|
"name": profile["name"],
|
|
"tile_id": tile_id,
|
|
"tile_count": "" if tile_count is None else str(tile_count),
|
|
"tile_area_fraction": "" if tile_area_fraction is None else f"{tile_area_fraction:.6f}",
|
|
"site_count": "" if site_count is None else str(site_count),
|
|
"request_weight": "" if weight is None else str(weight),
|
|
"wkt_chars": str(len(wkt)),
|
|
"status": str(response.get("status", "")),
|
|
"message": redact_sensitive_text(str(outputs.get("message", ""))),
|
|
"download_url": str(outputs.get("downloadUrl", "")),
|
|
"response_json": str(response_json),
|
|
"submitted_at_utc": datetime.now(timezone.utc).isoformat(timespec="seconds"),
|
|
}
|
|
|
|
|
|
def build_error_row(
|
|
county: dict[str, Any],
|
|
profile: dict[str, str],
|
|
site_count: int | None,
|
|
weight: int | None,
|
|
wkt: str,
|
|
error: Exception,
|
|
tile_id: str = "",
|
|
tile_count: int | None = None,
|
|
tile_area_fraction: float | None = None,
|
|
) -> dict[str, str]:
|
|
"""Build a structured error row."""
|
|
status_code = ""
|
|
request_url = ""
|
|
response = ""
|
|
if isinstance(error, NsrdArchiveRequestError):
|
|
status_code = str(error.status_code or "")
|
|
request_url = error.url or ""
|
|
response = error.response
|
|
|
|
return {
|
|
"county_fips": county["county_fips"],
|
|
"county_name": county["county_name"],
|
|
"state_fips": county["state_fips"],
|
|
"state_abbr": county["state_abbr"],
|
|
"profile_key": profile["profile_key"],
|
|
"source": profile["source"],
|
|
"tile_id": tile_id,
|
|
"tile_count": "" if tile_count is None else str(tile_count),
|
|
"tile_area_fraction": "" if tile_area_fraction is None else f"{tile_area_fraction:.6f}",
|
|
"site_count": "" if site_count is None else str(site_count),
|
|
"request_weight": "" if weight is None else str(weight),
|
|
"wkt_chars": str(len(wkt)),
|
|
"error_type": type(error).__name__,
|
|
"status_code": status_code,
|
|
"request_url": request_url,
|
|
"response": redact_sensitive_text(response),
|
|
"message": redact_sensitive_text(str(error)),
|
|
"logged_at_utc": datetime.now(timezone.utc).isoformat(timespec="seconds"),
|
|
}
|
|
|
|
|
|
def is_rate_limit_error(row: dict[str, str] | None) -> bool:
|
|
"""Return whether an error row represents a rate-limit stop condition."""
|
|
return row is not None and row.get("status_code") == "429"
|
|
|
|
|
|
def is_in_flight_limit_error(row: dict[str, str] | None) -> bool:
|
|
"""Return whether an error row means NSRDB has too many active archive jobs."""
|
|
if row is None or row.get("status_code") != "400":
|
|
return False
|
|
message = row.get("message", "").lower()
|
|
return "in-flight jobs" in message or "maximum limit of 20 requests" in message
|
|
|
|
|
|
def seed_queue_monitor_from_errors(
|
|
monitor: ArchiveQueueMonitor,
|
|
error_rows: dict[str, dict[str, str]],
|
|
) -> None:
|
|
"""Restore a recent server-reported queue count after a script restart."""
|
|
now = datetime.now(timezone.utc)
|
|
recent_rows: list[tuple[datetime, dict[str, str]]] = []
|
|
for row in error_rows.values():
|
|
if not is_in_flight_limit_error(row):
|
|
continue
|
|
logged_at = row.get("logged_at_utc", "").strip()
|
|
try:
|
|
logged = datetime.fromisoformat(logged_at.replace("Z", "+00:00"))
|
|
except ValueError:
|
|
continue
|
|
if logged.tzinfo is None:
|
|
logged = logged.replace(tzinfo=timezone.utc)
|
|
if (now - logged).total_seconds() <= DEFAULT_SERVER_COUNT_GRACE:
|
|
recent_rows.append((logged, row))
|
|
|
|
if recent_rows:
|
|
_, latest_row = max(recent_rows, key=lambda item: item[0])
|
|
monitor.observe_server_count(latest_row.get("message", ""))
|
|
|
|
|
|
def archive_download_status(url: str, timeout: int) -> str:
|
|
"""Return ready, pending, or unknown for an NSRDB archive URL."""
|
|
request = urllib.request.Request(
|
|
url,
|
|
headers={"User-Agent": "county-climate-explorer/0.1"},
|
|
method="HEAD",
|
|
)
|
|
try:
|
|
with urllib.request.urlopen(request, timeout=timeout):
|
|
return "ready"
|
|
except urllib.error.HTTPError as error:
|
|
host = urllib.parse.urlsplit(url).netloc.lower()
|
|
if error.code == 403 and (
|
|
host == "s3.amazonaws.com" or host.endswith(".amazonaws.com")
|
|
):
|
|
return "pending"
|
|
return "unknown"
|
|
except (OSError, urllib.error.URLError):
|
|
return "unknown"
|
|
|
|
|
|
def retry_after_from_message(message: str) -> str:
|
|
"""Extract a Retry-After value from the sanitized error message."""
|
|
match = re.search(r"Retry-After:\s*([^\s]+)", message)
|
|
return match.group(1) if match else ""
|
|
|
|
|
|
def wait_for_queue(cooldown: float, resume_message: str) -> None:
|
|
"""Pause for the NSRDB archive queue and announce when work resumes."""
|
|
if cooldown > 0:
|
|
time.sleep(cooldown)
|
|
print(resume_message)
|
|
|
|
|
|
def submit_geometry_request(
|
|
args: argparse.Namespace,
|
|
county: dict[str, Any],
|
|
profile: dict[str, str],
|
|
geometry: Any,
|
|
api_key: str,
|
|
email: str,
|
|
rate_limiter: RequestRateLimiter,
|
|
site_count: int | None = None,
|
|
weight: int | None = None,
|
|
tile_id: str = "",
|
|
tile_count: int | None = None,
|
|
tile_area_fraction: float | None = None,
|
|
) -> dict[str, str]:
|
|
"""Submit one county geometry or tile to the archive API."""
|
|
wkt = geometry_wkt(geometry, args.wkt_precision)
|
|
if site_count is None:
|
|
site_count = get_site_count(args, api_key, profile, wkt, rate_limiter)
|
|
if weight is None:
|
|
weight = request_weight(site_count, args.attributes, profile["name"], args.interval)
|
|
|
|
if site_count == 0:
|
|
raise RuntimeError(f"No {profile['source']} sites intersect this county polygon.")
|
|
if weight is not None and weight > args.max_request_weight:
|
|
raise ValueError(
|
|
f"Request weight {weight} exceeds max {args.max_request_weight}; tile this county before submitting."
|
|
)
|
|
|
|
response_json = response_path(args, county["county_fips"], profile["profile_key"], profile["name"], tile_id)
|
|
if args.dry_run:
|
|
response = {
|
|
"status": "dry-run",
|
|
"outputs": {
|
|
"message": "Dry run only; archive request was not submitted.",
|
|
"downloadUrl": "",
|
|
},
|
|
}
|
|
else:
|
|
tile_label = f" tile {tile_id}" if tile_id else ""
|
|
args.queue_monitor.wait_for_capacity(
|
|
f"{county['county_fips']}{tile_label}",
|
|
args.in_flight_retries,
|
|
)
|
|
response = post_form_json(
|
|
endpoint=profile["endpoint"],
|
|
api_key=api_key,
|
|
payload=archive_payload(args, email, profile, wkt),
|
|
timeout=args.timeout,
|
|
rate_limiter=rate_limiter,
|
|
)
|
|
|
|
errors = response.get("errors") or []
|
|
if errors:
|
|
raise NsrdArchiveRequestError(f"Archive API returned errors: {errors}")
|
|
|
|
args.response_dir.mkdir(parents=True, exist_ok=True)
|
|
response_json.write_text(json.dumps(response, indent=2, sort_keys=True), encoding="utf-8")
|
|
success_row = build_success_row(
|
|
county,
|
|
profile,
|
|
site_count,
|
|
weight,
|
|
wkt,
|
|
response,
|
|
response_json,
|
|
tile_id=tile_id,
|
|
tile_count=tile_count,
|
|
tile_area_fraction=tile_area_fraction,
|
|
)
|
|
if not args.dry_run:
|
|
args.queue_monitor.record_submission(success_row)
|
|
return success_row
|
|
|
|
|
|
def prepare_geometry_tiles(
|
|
args: argparse.Namespace,
|
|
county: dict[str, Any],
|
|
profile: dict[str, str],
|
|
api_key: str,
|
|
rate_limiter: RequestRateLimiter,
|
|
) -> list[dict[str, Any]]:
|
|
"""Split a county into requestable tiles with preflight site counts."""
|
|
if args.skip_site_count:
|
|
raise ValueError("Tiling requires site-count preflight; rerun without --skip-site-count.")
|
|
|
|
parent_geometry = county["geometry"]
|
|
pending: list[Any] = [parent_geometry]
|
|
accepted: list[dict[str, Any]] = []
|
|
|
|
while pending:
|
|
if len(pending) + len(accepted) > args.max_tiles_per_county:
|
|
raise ValueError(
|
|
f"Tiling {county['county_fips']} would exceed --max-tiles-per-county={args.max_tiles_per_county}."
|
|
)
|
|
|
|
geometry = clean_geometry(pending.pop(0))
|
|
if geometry is None or geometry.is_empty:
|
|
continue
|
|
|
|
wkt = geometry_wkt(geometry, args.wkt_precision)
|
|
should_pre_split = bool(args.tile_max_wkt_chars and len(wkt) > args.tile_max_wkt_chars)
|
|
if should_pre_split:
|
|
parts = split_geometry(geometry)
|
|
if len(parts) < 2:
|
|
raise ValueError(
|
|
f"Could not split oversized tile for {county['county_fips']} before site-count preflight."
|
|
)
|
|
pending.extend(parts)
|
|
continue
|
|
|
|
site_count = get_site_count(args, api_key, profile, wkt, rate_limiter)
|
|
if site_count is None:
|
|
parts = split_geometry(geometry)
|
|
if len(parts) < 2:
|
|
raise ValueError(
|
|
f"Could not split unchecked tile for {county['county_fips']} into site-countable pieces."
|
|
)
|
|
pending.extend(parts)
|
|
continue
|
|
weight = request_weight(site_count, args.attributes, profile["name"], args.interval)
|
|
if site_count == 0:
|
|
continue
|
|
|
|
if weight is not None and weight > args.max_request_weight:
|
|
parts = split_geometry(geometry)
|
|
if len(parts) < 2:
|
|
raise ValueError(
|
|
f"Could not split overweight tile for {county['county_fips']} with request weight {weight}."
|
|
)
|
|
pending.extend(parts)
|
|
continue
|
|
|
|
accepted.append(
|
|
{
|
|
"geometry": geometry,
|
|
"site_count": site_count,
|
|
"weight": weight,
|
|
"area_fraction": geometry_area_fraction(geometry, parent_geometry),
|
|
}
|
|
)
|
|
|
|
if not accepted:
|
|
raise RuntimeError(f"No {profile['source']} sites intersect this county polygon after tiling.")
|
|
return accepted
|
|
|
|
|
|
def submit_tiled_county_request(
|
|
args: argparse.Namespace,
|
|
county: dict[str, Any],
|
|
profile: dict[str, str],
|
|
api_key: str,
|
|
email: str,
|
|
rate_limiter: RequestRateLimiter,
|
|
) -> list[dict[str, str]]:
|
|
"""Submit all tiles for one county and return manifest rows."""
|
|
tiles = prepare_geometry_tiles(args, county, profile, api_key, rate_limiter)
|
|
tile_count = len(tiles)
|
|
rows: list[dict[str, str]] = []
|
|
for tile_index, tile in enumerate(tiles, start=1):
|
|
tile_id = f"tile{tile_index:03d}"
|
|
rows.append(
|
|
submit_geometry_request(
|
|
args,
|
|
county,
|
|
profile,
|
|
tile["geometry"],
|
|
api_key,
|
|
email,
|
|
rate_limiter,
|
|
site_count=tile["site_count"],
|
|
weight=tile["weight"],
|
|
tile_id=tile_id,
|
|
tile_count=tile_count,
|
|
tile_area_fraction=tile["area_fraction"],
|
|
)
|
|
)
|
|
return rows
|
|
|
|
|
|
def submit_county_request(
|
|
args: argparse.Namespace,
|
|
county: dict[str, Any],
|
|
api_key: str,
|
|
email: str,
|
|
rate_limiter: RequestRateLimiter,
|
|
) -> tuple[list[dict[str, str]], dict[str, str] | None]:
|
|
"""Submit the first valid archive profile for one county."""
|
|
wkt = county_wkt(county, args.wkt_precision)
|
|
final_error: Exception | None = None
|
|
final_profile: dict[str, str] | None = None
|
|
final_site_count: int | None = None
|
|
final_weight: int | None = None
|
|
|
|
if should_skip_independent_city(args, county):
|
|
profile = request_profiles(args, county)[0]
|
|
error = IndependentCitySkip(
|
|
"Skipped independent city with no intersecting NSRDB polygon grid sites; "
|
|
"use the representative-point metric instead."
|
|
)
|
|
return [], build_error_row(county, profile, None, None, wkt, error)
|
|
|
|
for profile in request_profiles(args, county):
|
|
site_count: int | None = None
|
|
weight: int | None = None
|
|
final_profile = profile
|
|
try:
|
|
return [
|
|
submit_geometry_request(args, county, profile, county["geometry"], api_key, email, rate_limiter)
|
|
], None
|
|
except (OSError, RuntimeError, ValueError, urllib.error.URLError) as error:
|
|
final_error = error
|
|
if isinstance(error, ValueError) and "exceeds max" in str(error) and args.tile_large_counties:
|
|
try:
|
|
return submit_tiled_county_request(args, county, profile, api_key, email, rate_limiter), None
|
|
except (OSError, RuntimeError, ValueError, urllib.error.URLError) as tile_error:
|
|
final_error = tile_error
|
|
break
|
|
if is_request_size_error(error) and args.tile_large_counties:
|
|
try:
|
|
return submit_tiled_county_request(args, county, profile, api_key, email, rate_limiter), None
|
|
except (OSError, RuntimeError, ValueError, urllib.error.URLError) as tile_error:
|
|
final_error = tile_error
|
|
break
|
|
if isinstance(error, RuntimeError) and "No " in str(error) and args.polar_fallback:
|
|
continue
|
|
break
|
|
|
|
if final_profile is None:
|
|
final_profile = request_profiles(args, county)[0]
|
|
if final_error is None:
|
|
final_error = RuntimeError("No NSRDB archive profile was available for this county.")
|
|
|
|
return [], build_error_row(county, final_profile, final_site_count, final_weight, wkt, final_error)
|
|
|
|
|
|
def run_batch(args: argparse.Namespace) -> None:
|
|
"""Submit and log county polygon archive requests."""
|
|
request_rows = read_existing_rows(args.requests_csv)
|
|
args.queue_monitor = ArchiveQueueMonitor(
|
|
args.requests_csv.parent,
|
|
DEFAULT_QUEUE_STATUS_TIMEOUT,
|
|
DEFAULT_QUEUE_STATUS_WORKERS,
|
|
)
|
|
|
|
touches_archive_api = not args.dry_run
|
|
touches_site_count_api = not args.skip_site_count
|
|
email = prompt_for_text("NLR/NREL API email: ", args.email) if touches_archive_api else (args.email or "")
|
|
api_key = (
|
|
prompt_for_secret("NLR/NREL API key: ", args.api_key)
|
|
if touches_archive_api or touches_site_count_api
|
|
else (args.api_key or "")
|
|
)
|
|
counties = slice_batch(load_counties(args.counties_geojson, args.include_puerto_rico), args.start, args.limit)
|
|
error_rows = read_existing_rows(args.error_csv)
|
|
seed_queue_monitor_from_errors(args.queue_monitor, error_rows)
|
|
rate_limiter = RequestRateLimiter(args.delay)
|
|
|
|
total = len(counties)
|
|
for index, county in enumerate(counties, start=1):
|
|
county_fips = county["county_fips"]
|
|
label = f"{county_fips} {county['county_name']}, {county['state_abbr']}"
|
|
machine = CountyRequestStateMachine(label)
|
|
existing_request_keys = county_row_keys(request_rows, county_fips)
|
|
existing_statuses = [request_rows[row_key].get("status") for row_key in existing_request_keys]
|
|
if existing_request_keys and any(status != "dry-run" for status in existing_statuses) and not args.overwrite:
|
|
machine.transition(CountyRequestEvent.EXISTING_REQUEST_FOUND)
|
|
print(f"[{index}/{total}] Skipped existing request for {label}")
|
|
continue
|
|
|
|
machine.transition(CountyRequestEvent.START)
|
|
in_flight_attempts = 0
|
|
while True:
|
|
successes, error = submit_county_request(args, county, api_key, email, rate_limiter)
|
|
if successes:
|
|
machine.transition(CountyRequestEvent.REQUEST_SUCCEEDED)
|
|
break
|
|
if not is_in_flight_limit_error(error):
|
|
break
|
|
|
|
args.queue_monitor.observe_server_count(error.get("message", ""))
|
|
error_rows[csv_row_key(error)] = error
|
|
write_rows(error_rows, args.error_csv, ERROR_FIELDS)
|
|
in_flight_attempts += 1
|
|
if in_flight_attempts > args.in_flight_retries:
|
|
machine.transition(CountyRequestEvent.REQUEST_FAILED)
|
|
print(
|
|
f"[{index}/{total}] NSRDB in-flight job limit is still active for {label} "
|
|
f"after {args.in_flight_retries} retries. Stopping batch."
|
|
)
|
|
break
|
|
|
|
machine.transition(CountyRequestEvent.RETRY_REQUIRED)
|
|
cooldown = args.queue_monitor.dynamic_wait_seconds()
|
|
print(
|
|
f"[{index}/{total}] NSRDB in-flight job limit reached for {label}; "
|
|
f"waiting {cooldown:.0f} seconds for the queue to catch up "
|
|
f"before retry {in_flight_attempts}/{args.in_flight_retries}."
|
|
)
|
|
args.queue_monitor.wait_for_dynamic_queue(
|
|
cooldown,
|
|
f"[{index}/{total}] Queue wait complete; resuming with {label}.",
|
|
)
|
|
machine.transition(CountyRequestEvent.RETRY_STARTED)
|
|
|
|
if machine.state is CountyRequestState.REQUESTED:
|
|
for success in successes:
|
|
request_rows[csv_row_key(success)] = success
|
|
for error_key in county_row_keys(error_rows, county_fips):
|
|
error_rows.pop(error_key, None)
|
|
write_rows(request_rows, args.requests_csv, REQUEST_FIELDS)
|
|
write_rows(error_rows, args.error_csv, ERROR_FIELDS)
|
|
tile_note = f", tiles={len(successes)}" if len(successes) > 1 else ""
|
|
print(
|
|
f"[{index}/{total}] Requested {label}: "
|
|
f"{successes[0]['profile_key']}, sites="
|
|
f"{sum(int(row['site_count']) for row in successes if row['site_count']) or 'unchecked'}"
|
|
f"{tile_note}"
|
|
)
|
|
continue
|
|
|
|
if error is not None:
|
|
if machine.state is CountyRequestState.SUBMITTING:
|
|
event = (
|
|
CountyRequestEvent.COUNTY_SKIPPED
|
|
if error.get("error_type") == "IndependentCitySkip"
|
|
else CountyRequestEvent.REQUEST_FAILED
|
|
)
|
|
machine.transition(event)
|
|
error_rows[csv_row_key(error)] = error
|
|
write_rows(request_rows, args.requests_csv, REQUEST_FIELDS)
|
|
write_rows(error_rows, args.error_csv, ERROR_FIELDS)
|
|
if machine.state is CountyRequestState.SKIPPED:
|
|
print(f"[{index}/{total}] Skipped {label}: {error['message']}")
|
|
continue
|
|
if error.get("error_type") == "LocalQueueCapacityError":
|
|
print(f"[{index}/{total}] {error['message']}")
|
|
break
|
|
print(f"[{index}/{total}] Failed {label}: {error['message']}")
|
|
if is_in_flight_limit_error(error):
|
|
print("Stopping batch because the NSRDB archive queue is still full after retrying.")
|
|
break
|
|
if is_rate_limit_error(error):
|
|
retry_after = retry_after_from_message(error["message"])
|
|
retry_note = f" Retry after {retry_after} seconds." if retry_after else ""
|
|
print(f"Stopping batch because the NSRDB API returned HTTP 429.{retry_note}")
|
|
break
|
|
|
|
print(f"Wrote {len(request_rows)} request rows to {args.requests_csv}")
|
|
print(f"Wrote {len(error_rows)} error rows to {args.error_csv}")
|
|
|
|
|
|
def parse_args(
|
|
argv: list[str] | None = None,
|
|
description: str | None = None,
|
|
) -> argparse.Namespace:
|
|
"""Parse command-line arguments."""
|
|
parser = argparse.ArgumentParser(description=description or __doc__)
|
|
parser.add_argument("--counties-geojson", type=Path, default=DEFAULT_COUNTIES_GEOJSON)
|
|
parser.add_argument("--requests-csv", type=Path)
|
|
parser.add_argument("--error-csv", type=Path)
|
|
parser.add_argument("--response-dir", type=Path)
|
|
parser.add_argument(
|
|
"--artifact-label",
|
|
help="Short metric label used in generated response JSON filenames.",
|
|
)
|
|
parser.add_argument(
|
|
"--legacy-artifact-label",
|
|
help="Optional previous filename label to reuse when a cached response already exists.",
|
|
)
|
|
parser.add_argument("--endpoint", default=DEFAULT_ENDPOINT)
|
|
parser.add_argument("--polar-endpoint", default=DEFAULT_POLAR_ENDPOINT)
|
|
parser.add_argument("--site-count-endpoint", default=DEFAULT_SITE_COUNT_ENDPOINT)
|
|
parser.add_argument("--profile", choices=["auto", "goes", "polar"], default="auto")
|
|
parser.add_argument("--name", default=DEFAULT_TMY_NAME, help="GOES TMY name, such as tmy or tmy-2024.")
|
|
parser.add_argument("--polar-name", default=DEFAULT_POLAR_TMY_NAME, help="Polar TMY name, usually tmy.")
|
|
parser.add_argument("--attributes", help="Comma-delimited NSRDB attributes to request.")
|
|
parser.add_argument("--interval", type=int, default=60, help="NSRDB interval in minutes.")
|
|
parser.add_argument("--utc", action="store_true", help="Request UTC timestamps instead of local standard time.")
|
|
parser.add_argument("--leap-day", action="store_true", help="Include leap day when available.")
|
|
parser.add_argument("--email", help="Email address registered with the NLR/NREL API.")
|
|
parser.add_argument("--api-key", help="API key. If omitted, the script prompts securely.")
|
|
parser.add_argument("--full-name", help="Optional requester full name.")
|
|
parser.add_argument("--affiliation", help="Optional requester affiliation.")
|
|
parser.add_argument("--reason", help="Optional request reason.")
|
|
parser.add_argument("--mailing-list", action="store_true", help="Opt in to the NSRDB mailing list.")
|
|
parser.add_argument("--start", type=int, default=0, help="Zero-based row offset after sorting counties by FIPS.")
|
|
parser.add_argument("--limit", type=int, default=5, help="Number of counties to request. Use --all for all counties.")
|
|
parser.add_argument("--all", action="store_true", help="Request all counties from --start onward.")
|
|
parser.add_argument("--delay", type=float, default=2.1, help="Minimum seconds between API requests.")
|
|
parser.add_argument(
|
|
"--in-flight-retries",
|
|
type=int,
|
|
default=DEFAULT_IN_FLIGHT_RETRIES,
|
|
help="Number of queue-full retries for the same county before stopping.",
|
|
)
|
|
parser.add_argument("--timeout", type=int, default=180, help="Request timeout in seconds.")
|
|
parser.add_argument("--wkt-precision", type=int, default=6, help="Decimal places for polygon WKT coordinates.")
|
|
parser.add_argument("--max-request-weight", type=int, default=DEFAULT_MAX_REQUEST_WEIGHT)
|
|
parser.add_argument("--skip-site-count", action="store_true", help="Submit without checking NSRDB site-count first.")
|
|
parser.add_argument(
|
|
"--no-tile-large-counties",
|
|
action="store_false",
|
|
dest="tile_large_counties",
|
|
help="Disable recursive tiling for oversized county polygon requests.",
|
|
)
|
|
parser.add_argument(
|
|
"--tile-max-wkt-chars",
|
|
type=int,
|
|
default=DEFAULT_TILE_MAX_WKT_CHARS,
|
|
help="Split tiles above this WKT character count before site-count preflight. Use 0 to disable.",
|
|
)
|
|
parser.add_argument(
|
|
"--max-tiles-per-county",
|
|
type=int,
|
|
default=DEFAULT_MAX_TILES_PER_COUNTY,
|
|
help="Maximum number of archive request tiles allowed for one county.",
|
|
)
|
|
parser.add_argument(
|
|
"--site-count-max-wkt-chars",
|
|
type=int,
|
|
default=DEFAULT_SITE_COUNT_MAX_WKT_CHARS,
|
|
help=(
|
|
"Skip the NSRDB site-count preflight for polygons with WKT longer than this many "
|
|
"characters. Use 0 to attempt site-count for every polygon."
|
|
),
|
|
)
|
|
parser.add_argument("--polar-min-latitude", type=float, default=DEFAULT_POLAR_MIN_LATITUDE)
|
|
parser.add_argument(
|
|
"--skip-independent-city-fips",
|
|
default=DEFAULT_SKIP_INDEPENDENT_CITY_FIPS,
|
|
help=(
|
|
"Comma-delimited independent-city FIPS values to mark as skipped because no polygon grid "
|
|
"sites intersect them. Use an empty string to disable."
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--no-polar-fallback",
|
|
action="store_false",
|
|
dest="polar_fallback",
|
|
help="Disable retrying high-latitude counties against the NSRDB Polar TMY endpoint.",
|
|
)
|
|
parser.add_argument("--include-puerto-rico", action="store_true", help="Include Puerto Rico county polygons.")
|
|
parser.add_argument("--overwrite", action="store_true", help="Resubmit counties already present in requests CSV.")
|
|
parser.add_argument("--dry-run", action="store_true", help="Build rows and site counts without submitting archive jobs.")
|
|
args = parser.parse_args(argv)
|
|
required_values = {
|
|
"--requests-csv": args.requests_csv,
|
|
"--error-csv": args.error_csv,
|
|
"--response-dir": args.response_dir,
|
|
"--artifact-label": args.artifact_label,
|
|
"--attributes": args.attributes,
|
|
}
|
|
missing = [option for option, value in required_values.items() if not value]
|
|
if missing:
|
|
parser.error(f"the following arguments are required: {', '.join(missing)}")
|
|
if args.delay < 0:
|
|
parser.error("--delay must be 0 or greater.")
|
|
if args.in_flight_retries < 0:
|
|
parser.error("--in-flight-retries must be 0 or greater.")
|
|
if args.interval <= 0:
|
|
parser.error("--interval must be greater than 0.")
|
|
if args.wkt_precision < 0:
|
|
parser.error("--wkt-precision must be 0 or greater.")
|
|
if args.tile_max_wkt_chars < 0:
|
|
parser.error("--tile-max-wkt-chars must be 0 or greater.")
|
|
if args.max_tiles_per_county < 1:
|
|
parser.error("--max-tiles-per-county must be 1 or greater.")
|
|
if args.site_count_max_wkt_chars < 0:
|
|
parser.error("--site-count-max-wkt-chars must be 0 or greater.")
|
|
if not re.fullmatch(r"[a-z0-9][a-z0-9_-]*", args.artifact_label):
|
|
parser.error("--artifact-label must contain only lowercase letters, digits, hyphens, or underscores.")
|
|
if args.legacy_artifact_label and not re.fullmatch(
|
|
r"[a-z0-9][a-z0-9_-]*",
|
|
args.legacy_artifact_label,
|
|
):
|
|
parser.error(
|
|
"--legacy-artifact-label must contain only lowercase letters, digits, hyphens, or underscores."
|
|
)
|
|
if args.all:
|
|
args.limit = None
|
|
return args
|
|
|
|
|
|
def main(
|
|
argv: list[str] | None = None,
|
|
description: str | None = None,
|
|
) -> None:
|
|
"""Run the NSRDB county polygon request workflow."""
|
|
args = parse_args(argv, description)
|
|
run_batch(args)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|