|
| 1 | +from __future__ import annotations |
| 2 | + |
| 3 | +import json |
| 4 | +import tempfile |
| 5 | +from dataclasses import dataclass |
| 6 | +from datetime import datetime, timezone |
| 7 | +from pathlib import Path |
| 8 | +from typing import Any, Mapping |
| 9 | + |
| 10 | +RUNTIME_REPORT_SCHEMA_VERSION = "runtime_report.v1" |
| 11 | + |
| 12 | + |
| 13 | +@dataclass(frozen=True) |
| 14 | +class RuntimeReportPersistResult: |
| 15 | + local_path: str | None = None |
| 16 | + gcs_uri: str | None = None |
| 17 | + |
| 18 | + |
| 19 | +def build_runtime_report_base( |
| 20 | + *, |
| 21 | + platform: str, |
| 22 | + deploy_target: str, |
| 23 | + service_name: str, |
| 24 | + strategy_profile: str, |
| 25 | + run_id: str, |
| 26 | + run_source: str, |
| 27 | + strategy_domain: str | None = None, |
| 28 | + account_scope: str | None = None, |
| 29 | + account_group: str | None = None, |
| 30 | + account_region: str | None = None, |
| 31 | + dry_run: bool = False, |
| 32 | + status: str = "started", |
| 33 | + started_at: datetime | str | None = None, |
| 34 | + finished_at: datetime | str | None = None, |
| 35 | + summary: Mapping[str, Any] | None = None, |
| 36 | + diagnostics: Mapping[str, Any] | None = None, |
| 37 | + artifacts: Mapping[str, Any] | None = None, |
| 38 | +) -> dict[str, Any]: |
| 39 | + return { |
| 40 | + "schema_version": RUNTIME_REPORT_SCHEMA_VERSION, |
| 41 | + "platform": str(platform), |
| 42 | + "deploy_target": str(deploy_target), |
| 43 | + "service_name": str(service_name), |
| 44 | + "strategy_profile": str(strategy_profile), |
| 45 | + "strategy_domain": _optional_string(strategy_domain), |
| 46 | + "account_scope": _resolve_account_scope( |
| 47 | + account_scope=account_scope, |
| 48 | + account_group=account_group, |
| 49 | + account_region=account_region, |
| 50 | + ), |
| 51 | + "account_group": _optional_string(account_group), |
| 52 | + "account_region": _optional_string(account_region), |
| 53 | + "run_id": str(run_id), |
| 54 | + "run_source": str(run_source), |
| 55 | + "status": str(status), |
| 56 | + "dry_run": bool(dry_run), |
| 57 | + "started_at": _normalize_datetime(started_at), |
| 58 | + "finished_at": _normalize_datetime(finished_at), |
| 59 | + "summary": _normalize_mapping(summary), |
| 60 | + "diagnostics": _normalize_mapping(diagnostics), |
| 61 | + "artifacts": _normalize_mapping(artifacts), |
| 62 | + "errors": [], |
| 63 | + } |
| 64 | + |
| 65 | + |
| 66 | +def finalize_runtime_report( |
| 67 | + report: dict[str, Any], |
| 68 | + *, |
| 69 | + status: str, |
| 70 | + finished_at: datetime | str | None = None, |
| 71 | + summary: Mapping[str, Any] | None = None, |
| 72 | + diagnostics: Mapping[str, Any] | None = None, |
| 73 | + artifacts: Mapping[str, Any] | None = None, |
| 74 | +) -> dict[str, Any]: |
| 75 | + report["status"] = str(status) |
| 76 | + report["finished_at"] = _normalize_datetime(finished_at or datetime.now(timezone.utc)) |
| 77 | + _merge_section(report, "summary", summary) |
| 78 | + _merge_section(report, "diagnostics", diagnostics) |
| 79 | + _merge_section(report, "artifacts", artifacts) |
| 80 | + return report |
| 81 | + |
| 82 | + |
| 83 | +def append_runtime_report_error( |
| 84 | + report: dict[str, Any], |
| 85 | + *, |
| 86 | + stage: str, |
| 87 | + message: str, |
| 88 | + **fields: Any, |
| 89 | +) -> dict[str, Any]: |
| 90 | + entry = { |
| 91 | + "stage": str(stage), |
| 92 | + "message": str(message), |
| 93 | + **_normalize_mapping(fields), |
| 94 | + } |
| 95 | + cleaned = _drop_empty(entry) |
| 96 | + report.setdefault("errors", []).append(cleaned) |
| 97 | + return cleaned |
| 98 | + |
| 99 | + |
| 100 | +def default_runtime_report_path( |
| 101 | + report: Mapping[str, Any], |
| 102 | + *, |
| 103 | + base_dir: str | Path | None = None, |
| 104 | +) -> Path: |
| 105 | + root = Path(base_dir).expanduser() if base_dir else Path(tempfile.gettempdir()) / "quant_runtime_reports" |
| 106 | + return root / runtime_report_relative_path(report) |
| 107 | + |
| 108 | + |
| 109 | +def runtime_report_relative_path(report: Mapping[str, Any]) -> Path: |
| 110 | + started_at = _coerce_datetime(report.get("started_at")) |
| 111 | + month_segment = started_at.strftime("%Y-%m") if started_at is not None else "unknown-month" |
| 112 | + segments = [ |
| 113 | + _sanitize_path_segment(report.get("platform")) or "unknown-platform", |
| 114 | + _sanitize_path_segment(report.get("strategy_profile")) or "unknown-profile", |
| 115 | + ] |
| 116 | + account_scope = _sanitize_path_segment(report.get("account_scope")) |
| 117 | + if account_scope: |
| 118 | + segments.append(account_scope) |
| 119 | + run_id = _sanitize_path_segment(report.get("run_id")) or "run" |
| 120 | + return Path(*segments, month_segment, f"{run_id}.json") |
| 121 | + |
| 122 | + |
| 123 | +def write_runtime_report_json( |
| 124 | + report: Mapping[str, Any], |
| 125 | + *, |
| 126 | + output_path: str | Path, |
| 127 | +) -> Path: |
| 128 | + path = Path(output_path) |
| 129 | + path.parent.mkdir(parents=True, exist_ok=True) |
| 130 | + payload = _normalize_mapping(report) |
| 131 | + path.write_text(json.dumps(payload, ensure_ascii=False, indent=2, sort_keys=True), encoding="utf-8") |
| 132 | + return path |
| 133 | + |
| 134 | + |
| 135 | +def build_runtime_report_gcs_uri( |
| 136 | + report: Mapping[str, Any], |
| 137 | + *, |
| 138 | + gcs_prefix_uri: str, |
| 139 | +) -> str: |
| 140 | + bucket_name, prefix = _parse_gcs_uri(gcs_prefix_uri) |
| 141 | + object_name = runtime_report_relative_path(report).as_posix() |
| 142 | + if prefix: |
| 143 | + object_name = f"{prefix.rstrip('/')}/{object_name}" |
| 144 | + return f"gs://{bucket_name}/{object_name}" |
| 145 | + |
| 146 | + |
| 147 | +def upload_runtime_report_to_gcs( |
| 148 | + report: Mapping[str, Any], |
| 149 | + *, |
| 150 | + gcs_uri: str, |
| 151 | + gcp_project_id: str | None = None, |
| 152 | + client_factory: Any = None, |
| 153 | +) -> str: |
| 154 | + bucket_name, object_name = _parse_gcs_uri(gcs_uri) |
| 155 | + if not object_name: |
| 156 | + raise ValueError(f"gcs_uri must include an object path, got: {gcs_uri!r}") |
| 157 | + if client_factory is None: |
| 158 | + try: |
| 159 | + from google.cloud import storage # type: ignore |
| 160 | + except ImportError as exc: |
| 161 | + raise RuntimeError("google-cloud-storage is required for GCS runtime report upload") from exc |
| 162 | + client_factory = storage.Client |
| 163 | + client = client_factory(project=gcp_project_id) if gcp_project_id is not None else client_factory() |
| 164 | + blob = client.bucket(bucket_name).blob(object_name) |
| 165 | + payload = json.dumps(_normalize_mapping(report), ensure_ascii=False, indent=2, sort_keys=True) |
| 166 | + blob.upload_from_string(payload, content_type="application/json") |
| 167 | + return f"gs://{bucket_name}/{object_name}" |
| 168 | + |
| 169 | + |
| 170 | +def persist_runtime_report( |
| 171 | + report: dict[str, Any], |
| 172 | + *, |
| 173 | + base_dir: str | Path | None = None, |
| 174 | + output_path: str | Path | None = None, |
| 175 | + gcs_prefix_uri: str | None = None, |
| 176 | + gcp_project_id: str | None = None, |
| 177 | + client_factory: Any = None, |
| 178 | +) -> RuntimeReportPersistResult: |
| 179 | + local_path = Path(output_path).expanduser() if output_path else default_runtime_report_path(report, base_dir=base_dir) |
| 180 | + gcs_uri = build_runtime_report_gcs_uri(report, gcs_prefix_uri=gcs_prefix_uri) if _optional_string(gcs_prefix_uri) else None |
| 181 | + _merge_section( |
| 182 | + report, |
| 183 | + "artifacts", |
| 184 | + { |
| 185 | + "runtime_report_local_path": str(local_path), |
| 186 | + }, |
| 187 | + ) |
| 188 | + write_runtime_report_json(report, output_path=local_path) |
| 189 | + if gcs_uri is not None: |
| 190 | + gcs_uri = upload_runtime_report_to_gcs( |
| 191 | + report, |
| 192 | + gcs_uri=gcs_uri, |
| 193 | + gcp_project_id=gcp_project_id, |
| 194 | + client_factory=client_factory, |
| 195 | + ) |
| 196 | + _merge_section( |
| 197 | + report, |
| 198 | + "artifacts", |
| 199 | + { |
| 200 | + "runtime_report_gcs_uri": gcs_uri, |
| 201 | + }, |
| 202 | + ) |
| 203 | + write_runtime_report_json(report, output_path=local_path) |
| 204 | + return RuntimeReportPersistResult(local_path=str(local_path), gcs_uri=gcs_uri) |
| 205 | + |
| 206 | + |
| 207 | +def _merge_section(report: dict[str, Any], key: str, payload: Mapping[str, Any] | None) -> None: |
| 208 | + if not payload: |
| 209 | + return |
| 210 | + current = dict(report.get(key) or {}) |
| 211 | + current.update(_normalize_mapping(payload)) |
| 212 | + report[key] = current |
| 213 | + |
| 214 | + |
| 215 | +def _resolve_account_scope( |
| 216 | + *, |
| 217 | + account_scope: str | None, |
| 218 | + account_group: str | None, |
| 219 | + account_region: str | None, |
| 220 | +) -> str | None: |
| 221 | + for value in (account_scope, account_group, account_region): |
| 222 | + normalized = _optional_string(value) |
| 223 | + if normalized is not None: |
| 224 | + return normalized |
| 225 | + return None |
| 226 | + |
| 227 | + |
| 228 | +def _optional_string(value: Any) -> str | None: |
| 229 | + if value is None: |
| 230 | + return None |
| 231 | + text = str(value).strip() |
| 232 | + return text or None |
| 233 | + |
| 234 | + |
| 235 | +def _normalize_datetime(value: datetime | str | None) -> str | None: |
| 236 | + coerced = _coerce_datetime(value) |
| 237 | + if coerced is None: |
| 238 | + return _optional_string(value) |
| 239 | + return coerced.astimezone(timezone.utc).isoformat().replace("+00:00", "Z") |
| 240 | + |
| 241 | + |
| 242 | +def _coerce_datetime(value: datetime | str | None) -> datetime | None: |
| 243 | + if isinstance(value, datetime): |
| 244 | + return value.astimezone(timezone.utc) |
| 245 | + text = _optional_string(value) |
| 246 | + if text is None: |
| 247 | + return None |
| 248 | + try: |
| 249 | + if text.endswith("Z"): |
| 250 | + return datetime.fromisoformat(text.replace("Z", "+00:00")).astimezone(timezone.utc) |
| 251 | + return datetime.fromisoformat(text).astimezone(timezone.utc) |
| 252 | + except ValueError: |
| 253 | + return None |
| 254 | + |
| 255 | + |
| 256 | +def _normalize_mapping(mapping: Mapping[str, Any] | None) -> dict[str, Any]: |
| 257 | + if not mapping: |
| 258 | + return {} |
| 259 | + return {str(key): _normalize_value(value) for key, value in mapping.items()} |
| 260 | + |
| 261 | + |
| 262 | +def _normalize_value(value: Any) -> Any: |
| 263 | + if isinstance(value, datetime): |
| 264 | + return value.astimezone(timezone.utc).isoformat().replace("+00:00", "Z") |
| 265 | + if isinstance(value, Path): |
| 266 | + return str(value) |
| 267 | + if isinstance(value, Mapping): |
| 268 | + return _drop_empty({str(key): _normalize_value(item) for key, item in value.items()}) |
| 269 | + if isinstance(value, tuple): |
| 270 | + return [_normalize_value(item) for item in value] |
| 271 | + if isinstance(value, list): |
| 272 | + return [_normalize_value(item) for item in value] |
| 273 | + return value |
| 274 | + |
| 275 | + |
| 276 | +def _drop_empty(payload: Mapping[str, Any]) -> dict[str, Any]: |
| 277 | + cleaned: dict[str, Any] = {} |
| 278 | + for key, value in payload.items(): |
| 279 | + if value is None: |
| 280 | + continue |
| 281 | + if isinstance(value, str) and not value.strip(): |
| 282 | + continue |
| 283 | + if isinstance(value, (list, tuple, dict)) and len(value) == 0: |
| 284 | + continue |
| 285 | + cleaned[str(key)] = value |
| 286 | + return cleaned |
| 287 | + |
| 288 | + |
| 289 | +def _sanitize_path_segment(value: Any) -> str | None: |
| 290 | + text = _optional_string(value) |
| 291 | + if text is None: |
| 292 | + return None |
| 293 | + safe = "".join(ch if ch.isalnum() or ch in {"-", "_", "."} else "_" for ch in text) |
| 294 | + return safe or None |
| 295 | + |
| 296 | + |
| 297 | +def _parse_gcs_uri(value: str) -> tuple[str, str]: |
| 298 | + text = _optional_string(value) |
| 299 | + if text is None or not text.startswith("gs://"): |
| 300 | + raise ValueError(f"Expected gs://bucket[/prefix] URI, got: {value!r}") |
| 301 | + remainder = text[5:] |
| 302 | + bucket_name, _, object_name = remainder.partition("/") |
| 303 | + bucket = _optional_string(bucket_name) |
| 304 | + if bucket is None: |
| 305 | + raise ValueError(f"GCS bucket name is missing in URI: {value!r}") |
| 306 | + return bucket, object_name.strip("/") |
0 commit comments