scale-forecasting 1.0.0__py3-none-any.whl

scale_forecasting/__init__.py

@@ -0,0 +1,275 @@
+"""scale-forecasting — massively-parallel time-series forecasting on Google Cloud.
+
+The package is deliberately flat and readable: one capability per file. Start with the
+architecture overview in the README, then read ``config.py`` (the run contract),
+``worker.py`` (the unit of work), and one file under ``models/`` to see the whole shape.
+
+Public surface — two ways in, one engine underneath:
+
+* **The easy path.** `Forecaster` — point it at a config and call
+  `dry_run` / `run`
+  / `review`. Thin wrapper over `run`; no logic of its
+  own.
+* **The manage path.** `Registry` — the operator surface over the registry a deployment writes
+  to: `doctor` (read-only health), `drop_run`, `sweep_orphans`, `snapshot`, `export`, `init`.
+  Same verbs as ``python -m scale_forecasting.registry.ops``.
+* **The direct path.** Drive Spark or Ray yourself and reuse the *same* model machinery:
+  `run_group` (pure, per-bucket), `make_group_runner` / `make_chunk_runner` (the
+  writer-attached ``applyInPandas`` / Ray-task closures), `chunk_cells`, and the unit of work
+  `run_cell`. Both paths run byte-identical cell code. See ``docs/using_the_sdk.md``.
+
+Import cost: ``import scale_forecasting`` is near-instant. Light names (config, settings, errors)
+are eager; the heavy names above load the model modules and are therefore **lazy** — resolved on
+first attribute access via PEP 562. Touching ``Forecaster``/``run``/``run_cell`` pays the model
+import cost; touching ``RunConfig``/``Settings`` does not. ``test_sdk.py`` guards this contract.
+"""
+
+from __future__ import annotations
+
+from importlib.metadata import PackageNotFoundError
+from importlib.metadata import version as _pkg_version
+from typing import TYPE_CHECKING
+
+# Eager (light): the config contract, infra identity, and the error hierarchy. None of these pull
+# the model modules, so they keep ``import scale_forecasting`` cheap.
+from .config import (
+    Fanout,
+    RunConfig,
+    Workload,
+    estimate_fanout,
+    estimate_workload,
+    load_config,
+)
+from .errors import (
+    ConfigError,
+    DataError,
+    EngineError,
+    ModelError,
+    RegistryError,
+    ScaleForecastError,
+)
+from .settings import Settings
+
+try:
+    __version__: str = _pkg_version("scale-forecasting")
+except PackageNotFoundError:  # pragma: no cover - zip-staged cloud worker without .dist-info
+    __version__ = "1.0.0"
+
+# Lazy (heavy): name -> (module, attribute). Resolved on first access by ``__getattr__`` (PEP 562)
+# so the model modules (statsmodels et al., ~5s) load only when a caller actually reaches for the
+# run/model surface — never merely to ``import scale_forecasting``.
+_LAZY: dict[str, tuple[str, str]] = {
+    # Easy path.
+    "Forecaster": (".sdk", "Forecaster"),
+    "Registry": (".sdk", "Registry"),
+    "DryRunResult": (".sdk", "DryRunResult"),
+    "RunResult": (".sdk", "RunResult"),
+    "ModelResult": (".sdk", "ModelResult"),
+    "JobTrace": (".sdk", "JobTrace"),
+    "build_explain_frame": (".sdk", "build_explain_frame"),
+    "build_best_params_frame": (".sdk", "build_best_params_frame"),
+    "build_trace_frame": (".sdk", "build_trace_frame"),
+    "plot_trace": (".sdk", "plot_trace"),
+    # Run inspection — monitor a running run, review a finished one.
+    "monitor_run": (".review", "monitor_run"),
+    "review_run": (".review", "review_run"),
+    "RunProgress": (".review", "RunProgress"),
+    "RunReview": (".review", "RunReview"),
+    "calibration_report": (".review", "calibration_report"),
+    "CalibrationReport": (".review", "CalibrationReport"),
+    "build_leaderboard_frame": (".review", "build_leaderboard_frame"),
+    "build_predictions_frame": (".review", "build_predictions_frame"),
+    "build_hierarchy_frame": (".review", "build_hierarchy_frame"),
+    "build_cohorts_frame": (".review", "build_cohorts_frame"),
+    "build_calibration_frames": (".review", "build_calibration_frames"),
+    "build_ensemble_weights_frame": (".review", "build_ensemble_weights_frame"),
+    "build_attributions_frame": (".review", "build_attributions_frame"),
+    "explain_forecast_frame": (".review", "explain_forecast_frame"),
+    "plot_progress": (".review", "plot_progress"),
+    "plot_leaderboard": (".review", "plot_leaderboard"),
+    "plot_metric_distribution": (".review", "plot_metric_distribution"),
+    "plot_forecasts_frame": (".review", "plot_forecasts_frame"),
+    "plot_hierarchy_frame": (".review", "plot_hierarchy_frame"),
+    "plot_calibration": (".review", "plot_calibration"),
+    "plot_ensemble_weights": (".review", "plot_ensemble_weights"),
+    "plot_attributions": (".review", "plot_attributions"),
+    "plot_forecast_explanation": (".review", "plot_forecast_explanation"),
+    "DagNode": (".dag", "DagNode"),
+    # Cross-run ensembling.
+    "merge_configs_for_ensemble": (".ensemble_run", "merge_configs_for_ensemble"),
+    "ensemble_cross_runs": (".ensemble_run", "ensemble_cross_runs"),
+    # Orchestration.
+    "run": (".main", "run"),
+    # Direct path — pure core + writer-attached runners + unit of work.
+    "run_group": (".engines.spark_io", "run_group"),
+    "make_group_runner": (".engines.spark_io", "make_group_runner"),
+    "RunOutcome": (".engines.spark_io", "RunOutcome"),
+    "aggregate_status": (".engines.spark_io", "aggregate_status"),
+    "make_chunk_runner": (".engines.ray_io", "make_chunk_runner"),
+    "chunk_cells": (".engines.ray_io", "chunk_cells"),
+    "run_cell": (".worker", "run_cell"),
+    "CellResult": (".worker", "CellResult"),
+    # Model registry + runtime routing + catalogs.
+    "get_model": (".models", "get_model"),
+    "list_models": (".models", "list_models"),
+    "model_catalog": (".playground", "model_catalog"),
+    "metric_catalog": (".playground", "metric_catalog"),
+    "split_by_runtime": (".router", "split_by_runtime"),
+    # Hierarchical aggregation + coherent reconciliation.
+    "HierarchySpec": (".reconciliation", "HierarchySpec"),
+    "build_hierarchy": (".reconciliation", "build_hierarchy"),
+    "reconcile_matrix": (".reconciliation", "reconcile_matrix"),
+    "reconcile_forecasts": (".reconciliation", "reconcile_forecasts"),
+    "reconcile_oof": (".reconciliation", "reconcile_oof"),
+    "reconcile_cells": (".reconciliation", "reconcile_cells"),
+    "verify_coherence": (".reconciliation", "verify_coherence"),
+    # AI agent surfaces & MCP server.
+    "probe_environment": (".agent_surfaces", "probe_environment"),
+    "McpServer": (".mcp", "McpServer"),
+    # Forecaster Workbench diagnostics, scenario simulation & self-contained reports.
+    "compute_demand_profile": (".features", "compute_demand_profile"),
+    "compute_fva": (".diagnostics", "compute_fva"),
+    "diebold_mariano_test": (".diagnostics", "diebold_mariano_test"),
+    "wilcoxon_pvalue": (".diagnostics", "wilcoxon_pvalue"),
+    "compute_horizon_decay": (".diagnostics", "compute_horizon_decay"),
+    "compute_bias_report": (".diagnostics", "compute_bias_report"),
+    "compute_realized_accuracy": (".diagnostics", "compute_realized_accuracy"),
+    "compute_revision_stability": (".diagnostics", "compute_revision_stability"),
+    "compute_exceptions_queue": (".diagnostics", "compute_exceptions_queue"),
+    "rollup_time": (".diagnostics", "rollup_time"),
+    "plot_fva": (".diagnostics", "plot_fva"),
+    "plot_horizon_decay": (".diagnostics", "plot_horizon_decay"),
+    "apply_scenario_shocks": (".scenarios", "apply_scenario_shocks"),
+    "simulate_scenarios": (".scenarios", "simulate_scenarios"),
+    "summarize_scenarios": (".scenarios", "summarize_scenarios"),
+    "plot_scenarios": (".scenarios", "plot_scenarios"),
+    "build_report_payload": (".report", "build_report_payload"),
+    "render_html_report": (".report", "render_html_report"),
+    "export_report": (".report", "export_report"),
+}
+
+__all__ = [
+    "__version__",
+    # Eager.
+    "RunConfig",
+    "load_config",
+    "estimate_workload",
+    "Workload",
+    "estimate_fanout",
+    "Fanout",
+    "Settings",
+    "ScaleForecastError",
+    "ConfigError",
+    "DataError",
+    "ModelError",
+    "RegistryError",
+    "EngineError",
+    # Lazy.
+    *_LAZY,
+]
+
+if TYPE_CHECKING:  # so IDEs / type-checkers see the lazy names as real imports (re-exported below)
+    from .agent_surfaces import probe_environment  # noqa: F401
+    from .dag import DagNode  # noqa: F401
+    from .diagnostics import (  # noqa: F401
+        compute_bias_report,
+        compute_exceptions_queue,
+        compute_fva,
+        compute_horizon_decay,
+        compute_realized_accuracy,
+        compute_revision_stability,
+        diebold_mariano_test,
+        plot_fva,
+        plot_horizon_decay,
+        rollup_time,
+        wilcoxon_pvalue,
+    )
+    from .engines.ray_io import chunk_cells, make_chunk_runner  # noqa: F401
+    from .engines.spark_io import (  # noqa: F401
+        RunOutcome,
+        aggregate_status,
+        make_group_runner,
+        run_group,
+    )
+    from .ensemble_run import ensemble_cross_runs, merge_configs_for_ensemble  # noqa: F401
+    from .features import compute_demand_profile  # noqa: F401
+    from .main import run  # noqa: F401
+    from .mcp import McpServer  # noqa: F401
+    from .models import get_model, list_models  # noqa: F401
+    from .playground import metric_catalog, model_catalog  # noqa: F401
+    from .reconciliation import (  # noqa: F401
+        HierarchySpec,
+        build_hierarchy,
+        reconcile_cells,
+        reconcile_forecasts,
+        reconcile_matrix,
+        reconcile_oof,
+        verify_coherence,
+    )
+    from .report import (  # noqa: F401
+        build_report_payload,
+        export_report,
+        render_html_report,
+    )
+    from .review import (  # noqa: F401
+        CalibrationReport,
+        RunProgress,
+        RunReview,
+        build_attributions_frame,
+        build_calibration_frames,
+        build_cohorts_frame,
+        build_ensemble_weights_frame,
+        build_hierarchy_frame,
+        build_leaderboard_frame,
+        build_predictions_frame,
+        calibration_report,
+        explain_forecast_frame,
+        monitor_run,
+        plot_attributions,
+        plot_calibration,
+        plot_ensemble_weights,
+        plot_forecast_explanation,
+        plot_forecasts_frame,
+        plot_hierarchy_frame,
+        plot_leaderboard,
+        plot_metric_distribution,
+        plot_progress,
+        review_run,
+    )
+    from .router import split_by_runtime  # noqa: F401
+    from .scenarios import (  # noqa: F401
+        apply_scenario_shocks,
+        plot_scenarios,
+        simulate_scenarios,
+        summarize_scenarios,
+    )
+    from .sdk import (  # noqa: F401
+        DryRunResult,
+        Forecaster,
+        JobTrace,
+        ModelResult,
+        Registry,
+        RunResult,
+        build_best_params_frame,
+        build_explain_frame,
+        build_trace_frame,
+        plot_trace,
+    )
+    from .worker import CellResult, run_cell  # noqa: F401
+
+
+def __getattr__(name: str) -> object:
+    """PEP 562 lazy loader: resolve a heavy public name on first access, then cache it."""
+    target = _LAZY.get(name)
+    if target is None:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    from importlib import import_module
+
+    module_name, attr = target
+    value = getattr(import_module(module_name, __name__), attr)
+    globals()[name] = value  # cache so subsequent lookups skip __getattr__
+    return value
+
+
+def __dir__() -> list[str]:
+    return sorted(__all__)

scale_forecasting/_entry.py

@@ -0,0 +1,110 @@
+"""Shared on-cluster entrypoint core for the Dataproc and Ray forecast launchers.
+
+Both `spark_entry` and `ray_entry` parse the same on-cluster contract — ``--config-uri`` (the run
+config, staged to GCS as JSON — the JSON *is* the reproducibility record) + the executed
+``--models`` subset + ``--manage-header`` + the ``--sf-*`` infra args, which `export_infra_env`
+promotes to ``os.environ`` before anything resolves `Settings` — then load the staged config and
+run the runtime's engine ``run(cfg, models=, manage_header=)``. This module is that shared skeleton,
+parametrized by a resolver so each launcher differs only in which engine callable it selects (Spark
+and Ray each have exactly one — there is no method switch).
+
+``--models m1,m2`` restricts the executed subset (the staged config's ``run_id`` is unchanged, so
+both runtimes share it) and ``--manage-header false`` puts the engine in contributor mode (``main``
+owns the single shared header). Both are optional — absent, the engine runs its standalone
+lifecycle over ``cfg.models``.
+
+``--provisioned-hardware gpu`` states what the submitter actually bought, so a model can select a
+device explicitly rather than asking its library to guess (see `hardware`). Emitted only by GPU
+jobs; absent everywhere else, which leaves device selection at ``auto`` — today's behaviour.
+
+The GCP/engine imports stay lazy so the launchers import cleanly offline.
+"""
+
+from __future__ import annotations
+
+import argparse
+from collections.abc import Callable
+from typing import TYPE_CHECKING
+
+from ._infra_args import add_infra_args, export_infra_env
+from .errors import configure_cli_logging, get_logger
+from .hardware import add_hardware_arg, export_hardware_env
+
+if TYPE_CHECKING:
+    from .config import RunConfig
+
+_log = get_logger(__name__)
+
+# A resolved (engine_run, label) pair: the callable to run and a short name for logging. The
+# callable takes ``(cfg, models=, manage_header=)`` — the shared engine contract.
+EngineResolver = Callable[[argparse.Namespace], "tuple[Callable[..., object], str]"]
+
+
+def parse_models(raw: str | None) -> list[str] | None:
+    """Parse the optional ``--models m1,m2`` CSV into a subset list (``None`` → run ``cfg.models``).
+
+    Empty/whitespace-only tokens are dropped so a trailing comma is harmless; an all-empty value
+    collapses to ``None`` (standalone) rather than an empty subset (which would run nothing).
+    """
+    if raw is None:
+        return None
+    names = [tok.strip() for tok in raw.split(",") if tok.strip()]
+    return names or None
+
+
+def build_parser(prog: str, description: str) -> argparse.ArgumentParser:
+    """The shared launcher parser: ``--config-uri`` + ``--models`` + ``--manage-header`` + infra."""
+    p = argparse.ArgumentParser(prog=prog, description=description)
+    p.add_argument("--config-uri", required=True, help="gs:// (or local) path to run config JSON")
+    p.add_argument(
+        "--models",
+        default=None,
+        help="optional comma-separated executed subset; absent runs all cfg.models",
+    )
+    p.add_argument(
+        "--manage-header",
+        default="true",
+        choices=("true", "false"),
+        help="false = contributor mode; main.run owns the shared header",
+    )
+    add_infra_args(p)
+    add_hardware_arg(p)
+    return p
+
+
+def run_entry(
+    argv: list[str] | None,
+    *,
+    prog: str,
+    description: str,
+    resolve_engine: EngineResolver,
+) -> None:
+    """Parse the shared launcher args, load the staged config, and run the resolved engine.
+
+    The one on-cluster driver skeleton for both runtimes: parse → export infra env → load the
+    staged config (`load_config_uri`, which reads a ``gs://`` URI directly or a local path) →
+    ``resolve_engine(ns)`` picks the engine callable → run it with the parsed subset + header mode.
+    """
+    from .config import load_config_uri
+
+    configure_cli_logging()
+    p = build_parser(prog, description)
+    ns = p.parse_args(argv)
+    export_infra_env(ns)
+    # The driver half of the device carrier. The executor half rides the job spec
+    # (`hardware.spark_executor_env` / `hardware.ray_env_vars`) because a driver's environment does
+    # not reach a Spark executor or a Ray task worker.
+    export_hardware_env(ns)
+    models = parse_models(ns.models)
+    manage_header = ns.manage_header == "true"
+    engine_run, label = resolve_engine(ns)
+    _log.info(
+        "%s: runtime=%s config_uri=%s models=%s manage_header=%s",
+        prog,
+        label,
+        ns.config_uri,
+        models,
+        manage_header,
+    )
+    cfg: RunConfig = load_config_uri(ns.config_uri)
+    engine_run(cfg, models=models, manage_header=manage_header)

scale_forecasting/_infra_args.py

@@ -0,0 +1,91 @@
+"""Infra-identity CLI args ↔ ``SF_*`` environment — the one delivery mapping.
+
+Dataproc Serverless allowlists Spark property prefixes and rejects driver-env
+(``spark.kubernetes.driverEnv.*`` → "unsupported properties"), so a batch cannot hand the
+infra identity to the driver as environment variables. Instead every entrypoint accepts the
+identity as ``--sf-*`` job args and exports them into ``os.environ`` before
+`resolve` runs — keeping env-based ``Settings`` the
+single local/cloud seam rather than forking a "resolve from args" path.
+
+This module owns that mapping in ONE place so the seed job, the Spark forecast launcher
+(``spark_entry``), the local submit helper (``submit``), and the Terraform modules that build the
+``--sf-*`` arg list all stay in agreement. Local runs pass no ``--sf-*`` and use the ambient
+``SF_*`` environment untouched.
+
+Public surface: ``INFRA_ARG_ENV``, ``add_infra_args``, ``export_infra_env``, ``infra_args_from``.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import argparse
+
+    from .settings import Settings
+
+# The (argparse dest → SF_* env var) mapping. ``add_infra_args`` registers ``--sf-*`` flags whose
+# argparse dests are these first elements; ``export_infra_env`` copies each provided value into the
+# matching env var before Settings.resolve().
+INFRA_ARG_ENV: tuple[tuple[str, str], ...] = (
+    ("sf_project_id", "SF_PROJECT_ID"),
+    ("sf_connection", "SF_CONNECTION"),
+    ("sf_warehouse_uri", "SF_WAREHOUSE_URI"),
+    ("sf_dataset_id", "SF_DATASET_ID"),
+    # Only emitted when the deployment actually splits the registry from the source panel (the
+    # value is "" otherwise and `infra_args_from` skips empty values), so an existing deployment's
+    # submitted command is byte-identical. It has to travel, though: without it a cluster driver
+    # would re-resolve to the *source* dataset and write the run's rows to the wrong registry.
+    ("sf_registry_dataset_id", "SF_REGISTRY_DATASET_ID"),
+    ("sf_region", "SF_REGION"),
+)
+
+
+def add_infra_args(parser: argparse.ArgumentParser) -> None:
+    """Register the ``--sf-*`` infra-identity flags on a parser (all default ``None``).
+
+    When unset (local runs) the ambient ``SF_*`` environment is used as-is; when set (a cluster
+    batch) `export_infra_env` promotes them to ``os.environ`` before resolution.
+    """
+    for dest, env_name in INFRA_ARG_ENV:
+        flag = "--" + dest.replace("_", "-")
+        parser.add_argument(
+            flag, type=str, default=None, help=f"sets {env_name} in the environment"
+        )
+
+
+def export_infra_env(ns: argparse.Namespace) -> None:
+    """Copy any provided ``--sf-*`` args from ``ns`` into ``os.environ`` (only when set).
+
+    Keeps env-based ``Settings.resolve()`` the single local/cloud seam without a parallel
+    "resolve from args" path. Values left unset are skipped, so a local run's ambient environment
+    is untouched.
+    """
+    for dest, env_name in INFRA_ARG_ENV:
+        value = getattr(ns, dest, None)
+        if value:
+            os.environ[env_name] = value
+
+
+def infra_args_from(settings: Settings) -> list[str]:
+    """Build the ``--sf-*`` arg list carrying ``settings`` to a cluster batch (submit side).
+
+    The inverse of `add_infra_args` / `export_infra_env`: turns a resolved
+    `Settings` into the flat ``["--sf-project-id", ...]`` list a
+    Dataproc batch passes so the driver re-materializes the same identity via env.
+    """
+    values = {
+        "sf_project_id": settings.project_id,
+        "sf_connection": settings.connection,
+        "sf_warehouse_uri": settings.warehouse_uri,
+        "sf_dataset_id": settings.dataset_id,
+        "sf_registry_dataset_id": settings.registry_dataset_id_override,
+        "sf_region": settings.region,
+    }
+    args: list[str] = []
+    for dest, _env_name in INFRA_ARG_ENV:
+        value = values[dest]
+        if value:
+            args += ["--" + dest.replace("_", "-"), value]
+    return args