Vertex AI AutoML & Tabular Workflows Engine

automl_engine

Vertex AI AutoML & Tabular Workflow for Forecasting execution engine.

Orchestrates global panel forecasting models on Vertex AI (runtime == "vertex_automl", family == "automl"):

• vertex_l2l: Vertex AI AutoML Forecasting (Learn-to-Learn NAS + Ensemble)
• vertex_tide: Time-series Dense Encoder (TiDE)
• vertex_tft: Temporal Fusion Transformer (TFT)
• vertex_seq2seq: Sequence-to-Sequence (Seq2Seq+)

Supports two execution modes configured via compute.automl_mode (or compute.families.automl.automl_mode):

1. "tabular_workflow" (default): Compiles and submits the transparent Kubeflow Pipelines (KFP v2) Tabular Workflow for Forecasting via google-cloud-pipeline-components and google.cloud.aiplatform.PipelineJob, with right-sized Dataflow evaluation workers, custom worker-pool machine/GPU overrides, automatic GCS extraction of stage_1_tuning_result_artifact_uri for zero-search warm-starting across runs (reuse_tuning_from_run_id), and BatchPredictionJob inference with explanations.
2. "training_job": Managed AutoMLForecastingTrainingJob / TimeSeriesDenseEncoderForecastingTrainingJob / TemporalFusionTransformerForecastingTrainingJob / SequenceToSequencePlusForecastingTrainingJob (see TRAINING_JOB_CLASSES) followed by BatchPredictionJob.

Every completed run writes standard CellResult records to forecast_metadata, forecast_predictions (including Tier 2 per-horizon-step explanations JSON), and backtest_oof, so AutoML models participate in the leaderboard, backtest coverage views, and downstream ensembling alongside Statistical, ML, Deep Learning, and BigQuery-native models.

AutoMLModelPlan dataclass

AutoMLModelPlan(
    model_type,
    automl_mode,
    root_dir,
    context_window,
    horizon,
    data_granularity_unit,
    data_granularity_count,
    optimization_objective,
    train_budget_milli_node_hours,
    quantiles,
    enable_explainability,
    stage_1_tuning_result_artifact_uri,
    reuse_tuning_from_run_id,
    max_num_trials,
    trainer_service_account,
    evaluation_dataflow_machine_type,
    evaluation_dataflow_starting_num_workers,
    evaluation_dataflow_max_num_workers,
    evaluation_batch_predict_machine_type,
    evaluation_batch_predict_starting_replica_count,
    evaluation_batch_predict_max_replica_count,
    stage_1_tuner_worker_pool_specs_override,
    stage_2_trainer_worker_pool_specs_override,
    column_specs,
    time_series_attribute_columns,
    available_at_forecast_columns,
    unavailable_at_forecast_columns,
    extra_params=dict(),
)

Resolved execution plan for one Vertex AI AutoML / Tabular Workflow model.

build_automl_column_specs

build_automl_column_specs(cfg, panel_df)

Classify panel columns into Vertex AI Forecasting covariate lists and transform specs (pure).

Returns (column_specs, time_series_attribute_columns, available_at_forecast_columns, unavailable_at_forecast_columns).

Source code in src/scale_forecasting/engines/automl_engine.py

def build_automl_column_specs(
    cfg: RunConfig, panel_df: pd.DataFrame
) -> tuple[list[dict[str, str]], list[str], list[str], list[str]]:
    """Classify panel columns into Vertex AI Forecasting covariate lists and transform specs (pure).

    Returns ``(column_specs, time_series_attribute_columns, available_at_forecast_columns,
    unavailable_at_forecast_columns)``.
    """
    ts_id_col = cfg.data.ts_id_col
    date_col = cfg.data.date_col
    target_col = cfg.data.target_col

    present_cols = set(panel_df.columns)
    static_cols = [c for c in cfg.features.static_covariates if c in present_cols]
    future_cols = [c for c in cfg.features.future_covariates if c in present_cols]
    past_cols = [c for c in cfg.features.past_covariates if c in present_cols]

    # Engineered calendar/Fourier features are deterministic functions of timestamp and therefore
    # known in advance at forecast time.
    reserved = {ts_id_col, date_col, target_col, _SPLIT_COL} | set(static_cols) | set(past_cols)
    engineered_future = [
        c
        for c in panel_df.columns
        if c not in reserved and c not in future_cols and (c.startswith(("cal_", "sin_", "cos_")))
    ]

    time_series_attribute_columns = list(dict.fromkeys(static_cols))
    available_at_forecast_columns = list(
        dict.fromkeys([date_col, *future_cols, *engineered_future])
    )
    unavailable_at_forecast_columns = list(dict.fromkeys([target_col, *past_cols]))

    all_feature_cols = list(
        dict.fromkeys(
            [
                *time_series_attribute_columns,
                *available_at_forecast_columns,
                *unavailable_at_forecast_columns,
            ]
        )
    )
    column_specs: list[dict[str, str]] = []
    for col in all_feature_cols:
        if col == date_col:
            dtype = "timestamp"
        elif col == target_col or pd.api.types.is_numeric_dtype(panel_df[col]):
            dtype = "numeric"
        else:
            dtype = "categorical"
        column_specs.append({"column_name": col, "data_type": dtype})

    return (
        column_specs,
        time_series_attribute_columns,
        available_at_forecast_columns,
        unavailable_at_forecast_columns,
    )

compile_tabular_workflow_spec

compile_tabular_workflow_spec(
    plan,
    *,
    project_id,
    region,
    train_bq_uri,
    ts_id_col,
    date_col,
    target_col,
    output_dir,
)

Compile the GCPC Tabular Workflow for Forecasting YAML and parameter dict.

Uses google_cloud_pipeline_components.preview.automl.forecasting pipeline builders: - vertex_l2l: get_learn_to_learn_forecasting_pipeline_and_parameters - vertex_tide: get_time_series_dense_encoder_forecasting_pipeline_and_parameters - vertex_tft: get_temporal_fusion_transformer_forecasting_pipeline_and_parameters - vertex_seq2seq: get_sequence_to_sequence_forecasting_pipeline_and_parameters

Source code in src/scale_forecasting/engines/automl_engine.py

def compile_tabular_workflow_spec(
    plan: AutoMLModelPlan,
    *,
    project_id: str,
    region: str,
    train_bq_uri: str,
    ts_id_col: str,
    date_col: str,
    target_col: str,
    output_dir: str | Path,
) -> tuple[str, dict[str, Any]]:
    """Compile the GCPC Tabular Workflow for Forecasting YAML and parameter dict.

    Uses ``google_cloud_pipeline_components.preview.automl.forecasting`` pipeline builders:
    - ``vertex_l2l``: ``get_learn_to_learn_forecasting_pipeline_and_parameters``
    - ``vertex_tide``: ``get_time_series_dense_encoder_forecasting_pipeline_and_parameters``
    - ``vertex_tft``: ``get_temporal_fusion_transformer_forecasting_pipeline_and_parameters``
    - ``vertex_seq2seq``: ``get_sequence_to_sequence_forecasting_pipeline_and_parameters``
    """
    import shutil

    from google_cloud_pipeline_components.preview.automl import forecasting as gcpc_fc

    out_dir = Path(output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    base_model = base_model_name(plan.model_type)

    # Pass {"auto": [...]} so Vertex AI Feature Transform Engine resolves BigQuery native column
    # types (TIMESTAMP, DATETIME, DATE, FLOAT64, STRING) directly from BigQuery table stats.
    transformations: dict[str, list[str]] = {
        "auto": [spec["column_name"] for spec in plan.column_specs]
    }
    run_eval = bool(plan.extra_params.get("run_evaluation", False))

    common_kwargs: dict[str, Any] = {
        "project": project_id,
        "location": region,
        "root_dir": plan.root_dir,
        "time_column": date_col,
        "time_series_identifier_columns": [ts_id_col],
        "target_column": target_col,
        "forecast_horizon": plan.horizon,
        "context_window": plan.context_window,
        "optimization_objective": plan.optimization_objective,
        "transformations": transformations,
        "train_budget_milli_node_hours": float(plan.train_budget_milli_node_hours),
        "time_series_attribute_columns": plan.time_series_attribute_columns,
        "available_at_forecast_columns": plan.available_at_forecast_columns,
        "unavailable_at_forecast_columns": plan.unavailable_at_forecast_columns,
        "data_source_bigquery_table_path": train_bq_uri,
        "predefined_split_key": _SPLIT_COL,
        "feature_transform_engine_dataflow_machine_type": plan.evaluation_dataflow_machine_type,
        "feature_transform_engine_dataflow_max_num_workers": (
            plan.evaluation_dataflow_max_num_workers
        ),
        "evaluation_dataflow_machine_type": plan.evaluation_dataflow_machine_type,
        "evaluation_dataflow_starting_num_workers": plan.evaluation_dataflow_starting_num_workers,
        "evaluation_dataflow_max_num_workers": plan.evaluation_dataflow_max_num_workers,
        "evaluation_batch_predict_machine_type": plan.evaluation_batch_predict_machine_type,
        "evaluation_batch_predict_starting_replica_count": (
            plan.evaluation_batch_predict_starting_replica_count
        ),
        "evaluation_batch_predict_max_replica_count": (
            plan.evaluation_batch_predict_max_replica_count
        ),
        "run_evaluation": run_eval,
    }
    if run_eval:
        bq_parts = train_bq_uri.removeprefix("bq://").split(".")
        if len(bq_parts) >= 2:
            common_kwargs["evaluated_examples_bigquery_path"] = f"bq://{bq_parts[0]}.{bq_parts[1]}"

    if plan.stage_1_tuning_result_artifact_uri:
        common_kwargs["stage_1_tuning_result_artifact_uri"] = (
            plan.stage_1_tuning_result_artifact_uri
        )
    if plan.stage_1_tuner_worker_pool_specs_override:
        common_kwargs["stage_1_tuner_worker_pool_specs_override"] = (
            plan.stage_1_tuner_worker_pool_specs_override
        )
    if plan.stage_2_trainer_worker_pool_specs_override:
        common_kwargs["stage_2_trainer_worker_pool_specs_override"] = (
            plan.stage_2_trainer_worker_pool_specs_override
        )
    if plan.trainer_service_account:
        common_kwargs["dataflow_service_account"] = plan.trainer_service_account
    if "dataflow_subnetwork" in plan.extra_params:
        common_kwargs["dataflow_subnetwork"] = str(plan.extra_params["dataflow_subnetwork"])
    if "dataflow_use_public_ips" in plan.extra_params:
        common_kwargs["dataflow_use_public_ips"] = bool(
            plan.extra_params["dataflow_use_public_ips"]
        )

    if "holiday_regions" in plan.extra_params:
        common_kwargs["holiday_regions"] = plan.extra_params["holiday_regions"]
    if "window_stride_length" in plan.extra_params:
        common_kwargs["window_stride_length"] = plan.extra_params["window_stride_length"]
    if "window_max_count" in plan.extra_params:
        common_kwargs["window_max_count"] = plan.extra_params["window_max_count"]
    if "stage_1_num_parallel_trials" in plan.extra_params:
        common_kwargs["stage_1_num_parallel_trials"] = int(
            plan.extra_params["stage_1_num_parallel_trials"]
        )
    elif "max_parallel_trial_count" in plan.extra_params:
        common_kwargs["stage_1_num_parallel_trials"] = int(
            plan.extra_params["max_parallel_trial_count"]
        )
    if "stage_2_num_parallel_trials" in plan.extra_params:
        common_kwargs["stage_2_num_parallel_trials"] = int(
            plan.extra_params["stage_2_num_parallel_trials"]
        )
    elif "stage_1_num_parallel_trials" in common_kwargs:
        common_kwargs["stage_2_num_parallel_trials"] = common_kwargs["stage_1_num_parallel_trials"]

    if base_model != "vertex_tft":
        if "num_selected_trials" in plan.extra_params:
            common_kwargs["num_selected_trials"] = int(plan.extra_params["num_selected_trials"])
        elif plan.max_num_trials < 10:
            common_kwargs["num_selected_trials"] = max(1, int(plan.max_num_trials))

    if base_model in ("vertex_l2l", "vertex_tide"):
        # Per Vertex AI Forecasting docs, enable_probabilistic_inference is incompatible with
        # minimize-quantile-loss; pass quantiles directly when using minimize-quantile-loss.
        if plan.optimization_objective == "minimize-quantile-loss":
            common_kwargs["enable_probabilistic_inference"] = False
            common_kwargs["quantiles"] = plan.quantiles
        elif plan.extra_params.get("enable_probabilistic_inference"):
            common_kwargs["enable_probabilistic_inference"] = True
        if "hierarchy_group_columns" in plan.extra_params:
            common_kwargs["group_columns"] = plan.extra_params["hierarchy_group_columns"]
        if "hierarchy_group_total_weight" in plan.extra_params:
            common_kwargs["group_total_weight"] = float(
                plan.extra_params["hierarchy_group_total_weight"]
            )
        if "hierarchy_temporal_total_weight" in plan.extra_params:
            common_kwargs["temporal_total_weight"] = float(
                plan.extra_params["hierarchy_temporal_total_weight"]
            )
        if "hierarchy_group_temporal_total_weight" in plan.extra_params:
            common_kwargs["group_temporal_total_weight"] = float(
                plan.extra_params["hierarchy_group_temporal_total_weight"]
            )

    if base_model == "vertex_tide":
        src_path, params = (
            gcpc_fc.get_time_series_dense_encoder_forecasting_pipeline_and_parameters(
                **common_kwargs
            )
        )
    elif base_model == "vertex_seq2seq":
        src_path, params = gcpc_fc.get_sequence_to_sequence_forecasting_pipeline_and_parameters(
            **common_kwargs
        )
    elif base_model == "vertex_tft":
        src_path, params = (
            gcpc_fc.get_temporal_fusion_transformer_forecasting_pipeline_and_parameters(
                **common_kwargs
            )
        )
    else:
        src_path, params = gcpc_fc.get_learn_to_learn_forecasting_pipeline_and_parameters(
            **common_kwargs
        )

    safe_model = plan.model_type.replace("@", "__")
    dest_path = out_dir / f"{safe_model}_pipeline.yaml"
    if Path(src_path).is_file() and Path(src_path).resolve() != dest_path.resolve():
        shutil.copyfile(src_path, dest_path)
        return str(dest_path), params
    return str(src_path), params

execute_automl_model_cells

execute_automl_model_cells(
    panel_df,
    model_type,
    cfg,
    *,
    run_id,
    settings,
    future_df=None,
    bq_client=None,
    fit_and_predict_fn=None,
    job_id=None,
)

Execute backtest folds (if enabled) and final forward forecast for one automl model.

Returns (cell_results, model_telemetry) where cell_results has one CellResult per series in panel_df with populated metrics, calibrated predictions, explanations, diagnostics["feature_attributions"], and best_params.

Source code in src/scale_forecasting/engines/automl_engine.py

def execute_automl_model_cells(
    panel_df: pd.DataFrame,
    model_type: str,
    cfg: RunConfig,
    *,
    run_id: str,
    settings: Settings,
    future_df: pd.DataFrame | None = None,
    bq_client: Any = None,
    fit_and_predict_fn: Any = None,
    job_id: str | None = None,
) -> tuple[list[CellResult], dict[str, Any]]:
    """Execute backtest folds (if enabled) and final forward forecast for one ``automl`` model.

    Returns ``(cell_results, model_telemetry)`` where ``cell_results`` has one `CellResult` per
    series in ``panel_df`` with populated metrics, calibrated predictions, ``explanations``,
    ``diagnostics["feature_attributions"]``, and ``best_params``.
    """
    t0 = time.perf_counter()
    cell_started_at = datetime.now(UTC)
    worker_id = _worker_id()
    dispatch = fit_and_predict_fn or _execute_automl_fit_and_predict
    eff_cfg = effective_config_for_model(cfg, get_model(model_type), model_type)
    ts_id_col = cfg.data.ts_id_col
    date_col = cfg.data.date_col
    target_col = cfg.data.target_col
    base_prefix = vertex_automl_job_id(job_id) if job_id else f"sf-{run_id[-8:]}"
    short_model = model_type.removeprefix("vertex_").replace("@", "-").replace("_", "-")
    m_period = seasonal_period(cfg.data.freq)

    df = panel_df.copy()
    df[ts_id_col] = df[ts_id_col].astype(str)
    df[date_col] = pd.to_datetime(df[date_col]).dt.tz_localize(None)
    df[target_col] = pd.to_numeric(df[target_col], errors="coerce").astype(float)
    df = df.sort_values([ts_id_col, date_col], kind="mergesort").reset_index(drop=True)

    series_groups: dict[str, pd.DataFrame] = {
        str(tid): grp.reset_index(drop=True) for tid, grp in df.groupby(ts_id_col, sort=False)
    }
    uids = sorted(series_groups.keys())

    authored = dict(cfg.model_params.get(model_type, {}))
    prior_tuning_uri: str | None = None
    if not authored.get("stage_1_tuning_result_artifact_uri") and authored.get(
        "reuse_tuning_from_run_id"
    ):
        prior_tuning_uri = resolve_prior_tuning_artifact_uri(
            str(authored["reuse_tuning_from_run_id"]),
            model_type,
            settings=settings,
            bq_client=bq_client,
        )

    artifact_root = settings.artifact_root
    default_sa = (
        getattr(settings, "compute_service_account", None)
        or os.environ.get("SF_COMPUTE_SA")
        or None
    )
    default_subnet = (
        getattr(settings, "subnetwork_uri", None) or os.environ.get("SF_SUBNETWORK_URI") or None
    )
    plan = plan_automl_model(
        cfg,
        model_type,
        df,
        run_id=run_id,
        horizon=cfg.data.horizon,
        artifact_root=artifact_root,
        service_account=default_sa,
        subnetwork_uri=default_subnet,
        prior_tuning_uri=prior_tuning_uri,
    )

    oof_by_uid: dict[str, pd.DataFrame | None] = dict.fromkeys(uids, None)
    metrics_by_uid: dict[str, dict[str, float]] = {
        uid: {name: float("nan") for name in METRIC_NAMES} for uid in uids
    }
    bt_status_by_uid: dict[str, str | None] = dict.fromkeys(uids, None)
    bt_note_by_uid: dict[str, str | None] = dict.fromkeys(uids, None)
    n_folds_by_uid: dict[str, int | None] = dict.fromkeys(uids, None)
    bt_refit_by_uid: dict[str, str | None] = dict.fromkeys(uids, None)
    ach_step_by_uid: dict[str, int | None] = dict.fromkeys(uids, None)
    ach_min_by_uid: dict[str, int | None] = dict.fromkeys(uids, None)
    first_val_by_uid: dict[str, date | None] = dict.fromkeys(uids, None)
    last_val_by_uid: dict[str, date | None] = dict.fromkeys(uids, None)
    n_fits = 0
    train_rows_total = 0

    # 1. Backtest folds (when cfg.backtest.enabled)
    if cfg.backtest.enabled:
        folds_by_uid = {uid: make_folds(len(series_groups[uid]), cfg) for uid in uids}
        max_folds = max((len(fl) for fl in folds_by_uid.values()), default=0)
        oof_chunks_by_uid: dict[str, list[pd.DataFrame]] = {u: [] for u in uids}
        fmetrics_by_uid: dict[str, list[dict[str, float]]] = {u: [] for u in uids}
        bt_horizon = cfg.backtest.gap + cfg.backtest.horizon

        for fold_idx in range(max_folds):
            train_end_map: dict[str, pd.Timestamp] = {}
            fold_lookup: dict[str, Any] = {}
            for uid in uids:
                folds = folds_by_uid[uid]
                if fold_idx < len(folds):
                    f = folds[fold_idx]
                    sub = series_groups[uid]
                    train_slice = sub.iloc[: f.train_end]
                    train_end_map[uid] = pd.Timestamp(train_slice[date_col].iloc[-1])
                    fold_lookup[uid] = f
                    train_rows_total += len(train_slice)

            if not train_end_map:
                continue

            active_df = df[df[ts_id_col].isin(train_end_map)].copy()
            fold_plan = plan_automl_model(
                cfg,
                model_type,
                active_df,
                run_id=run_id,
                horizon=bt_horizon,
                artifact_root=artifact_root,
                service_account=default_sa,
                subnetwork_uri=default_subnet,
                prior_tuning_uri=prior_tuning_uri,
            )
            fold_train = prepare_training_panel(
                active_df, eff_cfg, horizon=bt_horizon, train_end_by_series=train_end_map
            )
            fold_infer = prepare_inference_panel(
                active_df,
                eff_cfg,
                context_window=fold_plan.context_window,
                horizon=bt_horizon,
                origin_by_series=train_end_map,
            )
            fold_preds, _, _, fold_artifacts = dispatch(
                fold_plan,
                fold_train,
                fold_infer,
                eff_cfg,
                settings=settings,
                job_prefix=f"{base_prefix}-{short_model}-f{fold_idx + 1}",
                bq_client=bq_client,
            )
            n_fits += 1
            if prior_tuning_uri is None and fold_artifacts.get(
                "stage_1_tuning_result_artifact_uri"
            ):
                prior_tuning_uri = str(fold_artifacts["stage_1_tuning_result_artifact_uri"])

            for uid, f in fold_lookup.items():
                pred_df = fold_preds.get(uid)
                if pred_df is None or pred_df.empty:
                    continue
                sub = series_groups[uid]
                train_slice = sub.iloc[: f.train_end]
                val_slice = sub.iloc[f.val_start : f.val_end]
                val_dates = pd.to_datetime(val_slice[date_col]).dt.tz_localize(None)
                merged = pd.merge(
                    val_slice[[date_col, target_col]].assign(ds=val_dates),
                    pred_df,
                    on="ds",
                    how="inner",
                )
                if merged.empty:
                    continue
                y_train_raw = train_slice[target_col].to_numpy(dtype=float)
                y_val = merged[target_col].to_numpy(dtype=float)
                yhat_raw = merged["yhat_raw"].to_numpy(dtype=float)
                yhat_adjusted = merged["yhat"].to_numpy(dtype=float)
                yhat = yhat_raw if cfg.output.point_forecast == "raw" else yhat_adjusted
                lower = merged["yhat_lower"].to_numpy(dtype=float)
                upper = merged["yhat_upper"].to_numpy(dtype=float)
                fmetrics_by_uid[uid].append(
                    compute_metrics(
                        y_val,
                        yhat,
                        y_train=y_train_raw,
                        lower=lower,
                        upper=upper,
                        seasonal_period=m_period,
                    )
                )
                cutoff = pd.Timestamp(train_end_map[uid]).normalize()
                oof_chunks_by_uid[uid].append(
                    pd.DataFrame(
                        {
                            "ds": pd.to_datetime(merged["ds"]).to_numpy(),
                            "fold_id": f.fold_id,
                            "y_true": y_val,
                            "yhat": yhat,
                            "yhat_raw": yhat_raw,
                            "yhat_adjusted": yhat_adjusted,
                            "yhat_lower": lower,
                            "yhat_upper": upper,
                            "cutoff_date": cutoff,
                            "horizon_step": np.arange(1, len(y_val) + 1),
                            "yhat_stale": np.full(len(y_val), np.nan),
                        },
                        columns=list(OOF_COLUMNS),
                    )
                )

        for uid in uids:
            sub = series_groups[uid]
            flist = folds_by_uid[uid]
            n_ach = len(fmetrics_by_uid[uid])
            n_folds_by_uid[uid] = n_ach
            metrics_by_uid[uid] = rollup_metrics(fmetrics_by_uid[uid])
            bt_status_by_uid[uid], bt_note_by_uid[uid] = _backtest_outcome(n_ach, cfg, sub)
            if n_ach > 0:
                oof_by_uid[uid] = pd.concat(oof_chunks_by_uid[uid], ignore_index=True)
                bt_refit_by_uid[uid] = "per_fold"
                geom = resolve_geometry(len(sub), cfg)
                ach_step_by_uid[uid] = geom.step
                ach_min_by_uid[uid] = geom.min_train
                first_val_by_uid[uid] = pd.Timestamp(sub[date_col].iloc[flist[0].val_start]).date()
                last_val_by_uid[uid] = pd.Timestamp(
                    sub[date_col].iloc[flist[-1].val_end - 1]
                ).date()

    # 2. Final full-history fit + forward forecast
    if prior_tuning_uri and not plan.stage_1_tuning_result_artifact_uri:
        plan = plan_automl_model(
            cfg,
            model_type,
            df,
            run_id=run_id,
            horizon=cfg.data.horizon,
            artifact_root=artifact_root,
            service_account=default_sa,
            subnetwork_uri=default_subnet,
            prior_tuning_uri=prior_tuning_uri,
        )

    final_train = prepare_training_panel(df, eff_cfg, horizon=cfg.data.horizon)
    train_rows_total += len(final_train)
    final_infer = prepare_inference_panel(
        df,
        eff_cfg,
        context_window=plan.context_window,
        horizon=cfg.data.horizon,
        future_df=future_df,
    )
    preds_by_s, exps_by_s, global_attrs_by_s, artifacts_info = dispatch(
        plan,
        final_train,
        final_infer,
        eff_cfg,
        settings=settings,
        job_prefix=f"{base_prefix}-{short_model}-final",
        bq_client=bq_client,
    )
    n_fits += 1

    elapsed = max(0.0, time.perf_counter() - t0)
    per_series_fit_s = elapsed / max(1, len(uids))
    ended_at = datetime.now(UTC)

    best_params: dict[str, Any] = {
        "training_mode": "global",
        "automl_mode": plan.automl_mode,
        "context_window": plan.context_window,
        "optimization_objective": plan.optimization_objective,
        "train_budget_milli_node_hours": plan.train_budget_milli_node_hours,
    }
    if artifacts_info.get("stage_1_tuning_result_artifact_uri"):
        best_params["stage_1_tuning_result_artifact_uri"] = artifacts_info[
            "stage_1_tuning_result_artifact_uri"
        ]
    if artifacts_info.get("vertex_model_resource_name"):
        best_params["vertex_model_resource_name"] = artifacts_info["vertex_model_resource_name"]

    model_artifact_uri = (
        artifacts_info.get("stage_1_tuning_result_artifact_uri")
        or artifacts_info.get("vertex_model_resource_name")
        or plan.root_dir
    )

    metric = cfg.backtest.decision_metric
    corrected = corrected_arm_for(metric)

    results: list[CellResult] = []
    for idx, uid in enumerate(uids):
        s_df = series_groups[uid]
        raw_pred = preds_by_s.get(uid)
        if raw_pred is None or raw_pred.empty:
            fut_dates = _future_dates(
                pd.Timestamp(s_df[date_col].iloc[-1]), cfg.data.horizon, cfg.data.freq
            )
            raw_pred = pd.DataFrame(
                {
                    "ds": fut_dates,
                    "yhat": np.nan,
                    "yhat_raw": np.nan,
                    "yhat_lower": np.nan,
                    "yhat_upper": np.nan,
                    "quantiles": pd.array(["{}"] * len(fut_dates), dtype="string"),
                },
                columns=list(PREDICTION_COLUMNS),
            )

        oof = oof_by_uid[uid]
        arm, arm_decision = resolve_arm(cfg.output.point_forecast, oof, metric, corrected)
        cal = calibrate_from_oof(oof, DEFAULT_QUANTILES) if oof is not None else None
        calibrated, interval_calibration = apply_calibration(raw_pred, cal, arm)
        arm_comparison = compare_arms(oof, metric, corrected) if oof is not None else {}

        diag: dict[str, Any] = {
            "automl_mode": plan.automl_mode,
            "pipeline_root": plan.root_dir,
        }
        if uid in global_attrs_by_s:
            diag["feature_attributions"] = global_attrs_by_s[uid]
        if artifacts_info.get("evaluation_metrics_uri"):
            diag["evaluation_metrics_uri"] = artifacts_info["evaluation_metrics_uri"]
        if should_stamp_variant_diagnostics(model_type, cfg):
            diag.update(variant_diagnostics_for(model_type, cfg))

        results.append(
            CellResult(
                run_id=run_id,
                ts_id=uid,
                model_type=model_type,
                compute_engine="vertex_automl",
                model_hash=make_model_hash(run_id, uid, model_type, cfg),
                status="ok",
                error=None,
                predictions=calibrated,
                oof=oof,
                metrics=metrics_by_uid[uid],
                best_params=best_params,
                fit_seconds=per_series_fit_s,
                worker_id=worker_id,
                cell_started_at=cell_started_at,
                cell_ended_at=ended_at,
                model_artifact_uri=str(model_artifact_uri) if model_artifact_uri else None,
                backtest_status=bt_status_by_uid[uid],
                n_folds_achieved=n_folds_by_uid[uid],
                backtest_note=bt_note_by_uid[uid],
                backtest_refit=bt_refit_by_uid[uid],
                staleness_gap=None,
                achieved_step=ach_step_by_uid[uid],
                achieved_min_train=ach_min_by_uid[uid],
                first_val_date=first_val_by_uid[uid],
                last_val_date=last_val_by_uid[uid],
                interval_source="native",
                point_forecast_source=arm,
                point_forecast_decision=arm_decision,
                interval_calibration=interval_calibration,
                point_forecast_margin=arm_comparison.get("margin"),
                diagnostics=diag,
                explanations=exps_by_s.get(uid),
                n_fits=n_fits if idx == 0 else 0,
                train_rows_total=train_rows_total if idx == 0 else 0,
                n_hpo_fits=0,
            )
        )

    telemetry = {
        "model_type": model_type,
        "automl_mode": plan.automl_mode,
        "root_dir": plan.root_dir,
        "context_window": plan.context_window,
        "horizon": plan.horizon,
        "train_budget_milli_node_hours": plan.train_budget_milli_node_hours,
        **artifacts_info,
    }
    return results, telemetry

extract_pipeline_artifacts

extract_pipeline_artifacts(pipeline_job)

Extract tuning result GCS URI, uploaded Model resource name, and task telemetry (pure).

Inspects pipeline_job.gca_resource.job_detail.task_details from a completed Vertex AI Tabular Workflow PipelineJob to capture: - stage_1_tuning_result_artifact_uri (for warm-starting future runs) - vertex_model_resource_name (or unmanaged_container_model_uri for BatchPrediction) - evaluation_metrics_uri

Source code in src/scale_forecasting/engines/automl_engine.py

def extract_pipeline_artifacts(pipeline_job: Any) -> dict[str, Any]:
    """Extract tuning result GCS URI, uploaded Model resource name, and task telemetry (pure).

    Inspects ``pipeline_job.gca_resource.job_detail.task_details`` from a completed Vertex AI
    Tabular Workflow ``PipelineJob`` to capture:
    - ``stage_1_tuning_result_artifact_uri`` (for warm-starting future runs)
    - ``vertex_model_resource_name`` (or ``unmanaged_container_model_uri`` for BatchPrediction)
    - ``evaluation_metrics_uri``
    """
    info: dict[str, Any] = {
        "pipeline_job_resource_name": getattr(pipeline_job, "resource_name", None),
        "stage_1_tuning_result_artifact_uri": None,
        "vertex_model_resource_name": None,
        "unmanaged_container_model_uri": None,
        "evaluation_metrics_uri": None,
    }
    gca = getattr(pipeline_job, "gca_resource", None)
    job_detail = getattr(gca, "job_detail", None)
    task_details = getattr(job_detail, "task_details", None) or []

    for task in task_details:
        outputs = getattr(task, "outputs", None) or {}
        for out_key, out_val in outputs.items():
            artifacts = getattr(out_val, "artifacts", None) or []
            for art in artifacts:
                uri = getattr(art, "uri", "") or ""
                metadata = dict(getattr(art, "metadata", None) or {})
                if out_key == "tuning_result_output" and uri:
                    info["stage_1_tuning_result_artifact_uri"] = uri
                elif out_key == "model" and (uri or metadata.get("resourceName")):
                    info["vertex_model_resource_name"] = metadata.get("resourceName") or uri
                elif out_key == "unmanaged_container_model" and uri:
                    info["unmanaged_container_model_uri"] = uri
                elif out_key == "evaluation_metrics" and uri:
                    info["evaluation_metrics_uri"] = uri

    return info

map_freq_to_granularity

map_freq_to_granularity(freq)

Map a pandas frequency string (cfg.data.freq) to Vertex AI (unit, count) (pure).

Source code in src/scale_forecasting/engines/automl_engine.py

def map_freq_to_granularity(freq: str) -> tuple[str, int]:
    """Map a pandas frequency string (`cfg.data.freq`) to Vertex AI `(unit, count)` (pure)."""
    if freq in _GRANULARITY_MAP:
        return _GRANULARITY_MAP[freq]
    upper = freq.upper()
    if upper in _GRANULARITY_MAP:
        return _GRANULARITY_MAP[upper]
    if upper.startswith("W"):
        return ("week", 1)
    if upper.startswith(("M", "BM")):
        return ("month", 1)
    if upper.startswith("Q"):
        return ("month", 3)
    if upper.startswith(("Y", "A")):
        return ("year", 1)
    if upper.startswith("H"):
        return ("hour", 1)
    return ("day", 1)

parse_batch_predictions_and_explanations

parse_batch_predictions_and_explanations(
    bq_pred_df,
    *,
    ts_id_col,
    date_col,
    target_col,
    quantiles=None,
    intervals=None,
    expected_dates_by_series=None,
)

Parse Vertex AI BatchPredictionJob output table into predictions and explanations.

1. predictions_by_series[ts_id]: DataFrame with canonical PREDICTION_COLUMNS (ds, yhat, yhat_raw, yhat_lower, yhat_upper, quantiles).
2. explanations_by_series[ts_id]: Tier 2 per-horizon-step local explanation dicts.
3. global_attributions_by_series[ts_id]: Tier 1 series-level mean |attribution| per feature across the forecast horizon.

Source code in src/scale_forecasting/engines/automl_engine.py

def parse_batch_predictions_and_explanations(
    bq_pred_df: pd.DataFrame,
    *,
    ts_id_col: str,
    date_col: str,
    target_col: str,
    quantiles: list[float] | None = None,
    intervals: list[float] | None = None,
    expected_dates_by_series: dict[str, pd.DatetimeIndex] | None = None,
) -> tuple[
    dict[str, pd.DataFrame],
    dict[str, list[dict[str, Any] | None]],
    dict[str, dict[str, float]],
]:
    """Parse Vertex AI ``BatchPredictionJob`` output table into predictions and explanations.

    Returns:
    1. ``predictions_by_series[ts_id]``: DataFrame with canonical ``PREDICTION_COLUMNS``
       (``ds``, ``yhat``, ``yhat_raw``, ``yhat_lower``, ``yhat_upper``, ``quantiles``).
    2. ``explanations_by_series[ts_id]``: Tier 2 per-horizon-step local explanation dicts.
    3. ``global_attributions_by_series[ts_id]``: Tier 1 series-level mean ``|attribution|`` per
       feature across the forecast horizon.
    """
    eff_quantiles = quantiles or [0.1, 0.5, 0.9]
    eff_intervals = intervals or [0.8, 0.95]

    pred_col_candidates = [
        f"predicted_{target_col}",
        "predicted_on_target",
        "prediction",
        "yhat",
    ]
    pred_col = next((c for c in pred_col_candidates if c in bq_pred_df.columns), None)
    if pred_col is None:
        pred_col = next(
            (c for c in bq_pred_df.columns if c.startswith("predicted_")),
            target_col if target_col in bq_pred_df.columns else None,
        )
    if pred_col is None:
        raise EngineError(
            f"Vertex BatchPrediction output is missing a prediction column "
            f"(columns: {list(bq_pred_df.columns)})"
        )

    exp_col = next(
        (c for c in ("explanation", "explanations") if c in bq_pred_df.columns),
        None,
    )

    df = bq_pred_df.copy()
    df[ts_id_col] = df[ts_id_col].astype(str)
    df[date_col] = pd.to_datetime(df[date_col]).dt.tz_localize(None)
    if target_col in df.columns and target_col != pred_col:
        null_mask = df[target_col].isna()
        if null_mask.any():
            df = df[null_mask].copy()

    df = df.sort_values([ts_id_col, date_col], kind="mergesort").reset_index(drop=True)

    predictions_by_series: dict[str, pd.DataFrame] = {}
    explanations_by_series: dict[str, list[dict[str, Any] | None]] = {}
    global_attributions_by_series: dict[str, dict[str, float]] = {}

    for ts_id, grp in df.groupby(ts_id_col, sort=False):
        tid = str(ts_id)
        sub = grp
        if expected_dates_by_series is not None and tid in expected_dates_by_series:
            exp_dates = pd.to_datetime(expected_dates_by_series[tid]).tz_localize(None)
            sub = sub[sub[date_col].isin(exp_dates)]
        if sub.empty:
            continue

        ds_list: list[pd.Timestamp] = []
        yhat_list: list[float] = []
        lower_list: list[float] = []
        upper_list: list[float] = []
        quantiles_json_list: list[str] = []
        step_exps: list[dict[str, Any] | None] = []
        attr_accum: dict[str, list[float]] = {}

        for _, row in sub.iterrows():
            yhat, l80, u80, l95, u95 = _extract_prediction_value_and_bounds(
                row[pred_col], quantiles=eff_quantiles, intervals=eff_intervals
            )
            ds_list.append(pd.Timestamp(row[date_col]))
            yhat_list.append(yhat)
            lower_list.append(l80)
            upper_list.append(u80)
            q_dict = {
                k: v
                for k, v in (
                    ("0.025", l95),
                    ("0.1", l80),
                    ("0.5", yhat),
                    ("0.9", u80),
                    ("0.975", u95),
                )
                if math.isfinite(v)
            }
            quantiles_json_list.append(json.dumps(q_dict))

            step_exp = _extract_step_explanation(row[exp_col]) if exp_col else None
            step_exps.append(step_exp)
            if step_exp and isinstance(step_exp.get("attributions"), dict):
                for feat, val in step_exp["attributions"].items():
                    attr_accum.setdefault(feat, []).append(abs(float(val)))

        ds_ns = pd.DatetimeIndex(ds_list).as_unit("ns")
        pred_df = pd.DataFrame(
            {
                "ds": ds_ns,
                "yhat": np.asarray(yhat_list, dtype=float),
                "yhat_raw": np.asarray(yhat_list, dtype=float),
                "yhat_lower": np.asarray(lower_list, dtype=float),
                "yhat_upper": np.asarray(upper_list, dtype=float),
                "quantiles": pd.array(quantiles_json_list, dtype="string"),
            },
            columns=list(PREDICTION_COLUMNS),
        )
        predictions_by_series[tid] = pred_df
        if any(e is not None for e in step_exps):
            explanations_by_series[tid] = step_exps
        if attr_accum:
            global_attributions_by_series[tid] = {
                feat: round(float(np.mean(vals)), 6) for feat, vals in attr_accum.items()
            }

    return predictions_by_series, explanations_by_series, global_attributions_by_series

plan_automl_model

plan_automl_model(
    cfg,
    model_type,
    panel_df,
    *,
    run_id,
    horizon=None,
    artifact_root="gs://scale-forecasting-artifacts",
    service_account=None,
    subnetwork_uri=None,
    prior_tuning_uri=None,
)

Resolve the complete execution plan for one automl model (pure).

Source code in src/scale_forecasting/engines/automl_engine.py

def plan_automl_model(
    cfg: RunConfig,
    model_type: str,
    panel_df: pd.DataFrame,
    *,
    run_id: str,
    horizon: int | None = None,
    artifact_root: str = "gs://scale-forecasting-artifacts",
    service_account: str | None = None,
    subnetwork_uri: str | None = None,
    prior_tuning_uri: str | None = None,
) -> AutoMLModelPlan:
    """Resolve the complete execution plan for one ``automl`` model (pure)."""
    fc = cfg.resolve_family_compute("automl")
    automl_mode = fc.automl_mode or cfg.compute.automl_mode
    authored = dict(cfg.model_params.get(model_type, {}))
    eff_horizon = int(horizon or cfg.data.horizon)

    default_ctx = max(eff_horizon, min(eff_horizon * 2, 60))
    context_window = _as_int(authored.get("context_window", default_ctx), what="context_window")
    budget = _as_int(
        authored.get("train_budget_milli_node_hours", 1000), what="train_budget_milli_node_hours"
    )
    objective = str(authored.get("optimization_objective", "minimize-rmse"))
    quantiles = resolve_quantiles([0.8, 0.95], authored.get("quantiles"))
    enable_explainability = bool(
        authored.get("enable_explainability", authored.get("generate_explanation", True))
    )
    max_num_trials = _as_int(authored.get("max_num_trials", 15), what="max_num_trials")

    tuning_uri = (
        str(authored["stage_1_tuning_result_artifact_uri"])
        if authored.get("stage_1_tuning_result_artifact_uri")
        else prior_tuning_uri
    )
    reuse_run_id = (
        str(authored["reuse_tuning_from_run_id"])
        if authored.get("reuse_tuning_from_run_id")
        else None
    )

    unit, count = map_freq_to_granularity(cfg.data.freq)
    if "data_granularity_unit" in authored:
        unit = str(authored["data_granularity_unit"])
    if "data_granularity_count" in authored:
        count = _as_int(authored["data_granularity_count"], what="data_granularity_count")

    # Right-size Dataflow evaluation and BatchPrediction replica bounds so small/medium panels do
    # not default to 22-25 workers and never violate starting <= max.
    max_df_workers = _as_int(
        authored.get("evaluation_dataflow_max_num_workers")
        or authored.get("dataflow_max_num_workers")
        or fc.max_workers
        or fc.workers
        or 2,
        what="evaluation_dataflow_max_num_workers",
    )
    start_df_workers = min(
        _as_int(
            authored.get("evaluation_dataflow_starting_num_workers") or fc.min_workers or 1,
            what="evaluation_dataflow_starting_num_workers",
        ),
        max_df_workers,
    )
    df_machine = str(
        authored.get("evaluation_dataflow_machine_type")
        or authored.get("dataflow_machine_type")
        or "n1-standard-4"
    )

    max_bp_replicas = _as_int(
        authored.get("evaluation_batch_predict_max_replica_count")
        or fc.max_workers
        or fc.workers
        or 2,
        what="evaluation_batch_predict_max_replica_count",
    )
    start_bp_replicas = min(
        _as_int(
            authored.get("evaluation_batch_predict_starting_replica_count") or fc.min_workers or 1,
            what="evaluation_batch_predict_starting_replica_count",
        ),
        max_bp_replicas,
    )
    bp_machine = str(authored.get("evaluation_batch_predict_machine_type", "n1-highmem-8"))

    # Worker pool overrides for Stage 1 Tuner and Stage 2 Trainer when hardware=="gpu" or explicit
    # machine_type is configured.
    tuner_override: list[dict[str, Any]] | None = None
    trainer_override: list[dict[str, Any]] | None = None
    if (
        fc.hardware == "gpu"
        or fc.machine_type is not None
        or "trainer_machine_type" in authored
        or "trainer_replica_count" in authored
    ):
        mt = str(authored.get("trainer_machine_type") or fc.machine_type or "n1-standard-8")
        # Vertex AutoML Stage 1/2 custom training containers support n1-* / g2-* / a2-* shapes; map
        # n2-standard-* to n1-standard-* when overriding Tabular Workflow worker pools.
        if mt.startswith("n2-standard-"):
            mt = mt.replace("n2-standard-", "n1-standard-", 1)
        machine_spec: dict[str, Any] = {"machine_type": mt}
        if fc.hardware == "gpu" and fc.gpu_type:
            machine_spec["accelerator_type"] = _VERTEX_ACCELERATOR_MAP.get(
                fc.gpu_type, "NVIDIA_TESLA_T4"
            )
            machine_spec["accelerator_count"] = max(1, fc.accelerator_count)
        replicas = max(
            1,
            _as_int(
                authored.get("trainer_replica_count") or fc.workers or 1,
                what="trainer_replica_count",
            ),
        )
        pool_spec = [{"machine_spec": machine_spec, "replica_count": replicas}]
        tuner_override = pool_spec
        trainer_override = pool_spec

    eff_cfg = effective_config_for_model(cfg, get_model(model_type), model_type)
    (
        col_specs,
        attr_cols,
        avail_cols,
        unavail_cols,
    ) = build_automl_column_specs(eff_cfg, panel_df)

    extra: dict[str, Any] = {}
    for k in (
        "holiday_regions",
        "hierarchy_group_columns",
        "hierarchy_group_total_weight",
        "hierarchy_temporal_total_weight",
        "hierarchy_group_temporal_total_weight",
        "window_stride_length",
        "window_max_count",
        "additional_experiments",
        "run_evaluation",
        "enable_probabilistic_inference",
        "max_parallel_trial_count",
        "stage_1_num_parallel_trials",
        "stage_2_num_parallel_trials",
        "num_selected_trials",
        "dataflow_disk_size_gb",
    ):
        if k in authored:
            extra[k] = authored[k]

    df_subnetwork = authored.get("dataflow_subnetwork") or subnetwork_uri
    if df_subnetwork:
        extra["dataflow_subnetwork"] = str(df_subnetwork)
        extra["dataflow_use_public_ips"] = bool(authored.get("dataflow_use_public_ips", False))
    elif "dataflow_use_public_ips" in authored:
        extra["dataflow_use_public_ips"] = bool(authored["dataflow_use_public_ips"])

    resolved_sa = (
        str(authored.get("trainer_service_account") or authored.get("dataflow_service_account"))
        if (authored.get("trainer_service_account") or authored.get("dataflow_service_account"))
        else service_account
    )

    clean_root = artifact_root.rstrip("/")
    safe_model = model_type.replace("@", "__")
    root_dir = f"{clean_root}/{run_id}/tabular_workflow/{safe_model}"

    return AutoMLModelPlan(
        model_type=model_type,
        automl_mode=automl_mode,
        root_dir=root_dir,
        context_window=context_window,
        horizon=eff_horizon,
        data_granularity_unit=unit,
        data_granularity_count=count,
        optimization_objective=objective,
        train_budget_milli_node_hours=budget,
        quantiles=quantiles,
        enable_explainability=enable_explainability,
        stage_1_tuning_result_artifact_uri=tuning_uri,
        reuse_tuning_from_run_id=reuse_run_id,
        max_num_trials=max_num_trials,
        trainer_service_account=resolved_sa,
        evaluation_dataflow_machine_type=df_machine,
        evaluation_dataflow_starting_num_workers=start_df_workers,
        evaluation_dataflow_max_num_workers=max_df_workers,
        evaluation_batch_predict_machine_type=bp_machine,
        evaluation_batch_predict_starting_replica_count=start_bp_replicas,
        evaluation_batch_predict_max_replica_count=max_bp_replicas,
        stage_1_tuner_worker_pool_specs_override=tuner_override,
        stage_2_trainer_worker_pool_specs_override=trainer_override,
        column_specs=col_specs,
        time_series_attribute_columns=attr_cols,
        available_at_forecast_columns=avail_cols,
        unavailable_at_forecast_columns=unavail_cols,
        extra_params=extra,
    )

prepare_inference_panel

prepare_inference_panel(
    panel_df,
    cfg,
    *,
    context_window,
    horizon,
    origin_by_series=None,
    future_df=None,
)

Build the Vertex AI BatchPredictionJob input table (context + future NULL-target rows).

Vertex AI Forecasting BatchPredictionJob requires each series to include its historical context rows leading up to the forecast origin followed immediately by horizon future rows where target_col is NULL (NaN) and available_at_forecast_columns (timestamp + future covariates) are populated.

Source code in src/scale_forecasting/engines/automl_engine.py

def prepare_inference_panel(
    panel_df: pd.DataFrame,
    cfg: RunConfig,
    *,
    context_window: int,
    horizon: int,
    origin_by_series: dict[str, pd.Timestamp] | None = None,
    future_df: pd.DataFrame | None = None,
) -> pd.DataFrame:
    """Build the Vertex AI ``BatchPredictionJob`` input table (context + future NULL-target rows).

    Vertex AI Forecasting ``BatchPredictionJob`` requires each series to include its historical
    context rows leading up to the forecast origin followed immediately by ``horizon`` future rows
    where ``target_col`` is ``NULL`` (`NaN`) and ``available_at_forecast_columns`` (timestamp +
    future covariates) are populated.
    """
    ts_id_col = cfg.data.ts_id_col
    date_col = cfg.data.date_col
    target_col = cfg.data.target_col

    df = panel_df.copy()
    df[date_col] = pd.to_datetime(df[date_col])
    df[target_col] = pd.to_numeric(df[target_col], errors="coerce").astype(float)
    df = df.sort_values([ts_id_col, date_col], kind="mergesort").reset_index(drop=True)

    future_lookup: dict[str, pd.DataFrame] = {}
    if future_df is not None and not future_df.empty:
        fdf = future_df.copy()
        fdf[date_col] = pd.to_datetime(fdf[date_col])
        for tid, fgrp in fdf.groupby(ts_id_col, sort=False):
            future_lookup[str(tid)] = fgrp.sort_values(date_col, kind="mergesort")

    static_cols = [c for c in cfg.features.static_covariates if c in df.columns]
    future_cols = [c for c in cfg.features.future_covariates if c in df.columns]
    past_cols = [c for c in cfg.features.past_covariates if c in df.columns]

    keep_context = max(context_window, horizon * 2, 10)
    frames: list[pd.DataFrame] = []

    for ts_id, grp in df.groupby(ts_id_col, sort=False):
        tid = str(ts_id)
        if origin_by_series is not None and tid in origin_by_series:
            origin = pd.Timestamp(origin_by_series[tid])
            hist = (
                grp[grp[date_col] <= origin].dropna(subset=[target_col]).tail(keep_context).copy()
            )
            held_out = grp[grp[date_col] > origin].head(horizon).copy()
        else:
            hist = grp.dropna(subset=[target_col]).tail(keep_context).copy()
            held_out = pd.DataFrame()

        if hist.empty:
            continue

        last_date = pd.Timestamp(hist[date_col].iloc[-1])
        fut_dates = _future_dates(last_date, horizon, cfg.data.freq)

        fut = pd.DataFrame({ts_id_col: [tid] * horizon, date_col: fut_dates})
        fut[target_col] = np.nan

        # Populate static covariates from the last historical row.
        for scol in static_cols:
            fut[scol] = hist[scol].iloc[-1]

        # Populate future covariates from future_df, held-out fold rows, or forward-fill.
        for fcol in future_cols:
            if tid in future_lookup and fcol in future_lookup[tid].columns:
                merged = pd.merge(
                    fut[[date_col]],
                    future_lookup[tid][[date_col, fcol]].drop_duplicates(subset=[date_col]),
                    on=date_col,
                    how="left",
                )
                fut[fcol] = merged[fcol].ffill().bfill().fillna(hist[fcol].iloc[-1]).to_numpy()
            elif not held_out.empty and fcol in held_out.columns:
                merged = pd.merge(
                    fut[[date_col]],
                    held_out[[date_col, fcol]].drop_duplicates(subset=[date_col]),
                    on=date_col,
                    how="left",
                )
                fut[fcol] = merged[fcol].ffill().bfill().fillna(hist[fcol].iloc[-1]).to_numpy()
            else:
                fut[fcol] = hist[fcol].iloc[-1]

        # Past covariates are unavailable in the future horizon -> NULL.
        for pcol in past_cols:
            fut[pcol] = np.nan

        combined = pd.concat([hist, fut], ignore_index=True)
        if _SPLIT_COL in combined.columns:
            combined = combined.drop(columns=[_SPLIT_COL])
        frames.append(combined)

    if not frames:
        raise EngineError("prepare_inference_panel produced 0 rows across all series")
    out = pd.concat(frames, ignore_index=True)
    out[ts_id_col] = out[ts_id_col].astype(str)
    for col in out.columns:
        if col not in (ts_id_col, date_col, _SPLIT_COL) and pd.api.types.is_numeric_dtype(df[col]):
            out[col] = pd.to_numeric(out[col], errors="coerce").astype(float)
    return out

prepare_training_panel

prepare_training_panel(
    panel_df, cfg, *, horizon=None, train_end_by_series=None
)

Prepare the BigQuery training table DataFrame with deterministic __sf_split__ (pure).

Slices each series to <= train_end_by_series[ts_id] when provided (backtest fold), sorts chronologically, and assigns TRAIN, VALIDATE, and TEST splits per series so Vertex AI Tabular Workflow's predefined_split_key receives a leak-free temporal split.

Source code in src/scale_forecasting/engines/automl_engine.py

def prepare_training_panel(
    panel_df: pd.DataFrame,
    cfg: RunConfig,
    *,
    horizon: int | None = None,
    train_end_by_series: dict[str, pd.Timestamp] | None = None,
) -> pd.DataFrame:
    """Prepare the BigQuery training table DataFrame with deterministic ``__sf_split__`` (pure).

    Slices each series to ``<= train_end_by_series[ts_id]`` when provided (backtest fold), sorts
    chronologically, and assigns ``TRAIN``, ``VALIDATE``, and ``TEST`` splits per series so Vertex
    AI Tabular Workflow's ``predefined_split_key`` receives a leak-free temporal split.
    """
    ts_id_col = cfg.data.ts_id_col
    date_col = cfg.data.date_col
    target_col = cfg.data.target_col
    eff_horizon = horizon or cfg.data.horizon

    df = panel_df.copy()
    df[date_col] = pd.to_datetime(df[date_col])
    df[target_col] = pd.to_numeric(df[target_col], errors="coerce").astype(float)
    df = df.sort_values([ts_id_col, date_col], kind="mergesort").reset_index(drop=True)

    frames: list[pd.DataFrame] = []
    for ts_id, grp in df.groupby(ts_id_col, sort=False):
        sub = grp
        if train_end_by_series is not None and str(ts_id) in train_end_by_series:
            cutoff = pd.Timestamp(train_end_by_series[str(ts_id)])
            sub = sub[sub[date_col] <= cutoff]
        sub = sub.dropna(subset=[target_col]).copy()
        n = len(sub)
        if n == 0:
            continue
        # Ensure VALIDATE and TEST each hold at least `min(eff_horizon, max(1, n // 5))` rows while
        # TRAIN holds at least 60% of the series history.
        eval_win = min(eff_horizon, max(1, n // 5))
        if n >= 3:
            test_start = n - eval_win
            val_start = max(1, test_start - eval_win)
            splits = (
                ["TRAIN"] * val_start
                + ["VALIDATE"] * (test_start - val_start)
                + ["TEST"] * (n - test_start)
            )
        else:
            splits = ["TRAIN"] * n
        sub[_SPLIT_COL] = splits
        frames.append(sub)

    if not frames:
        raise EngineError("prepare_training_panel produced 0 rows across all series")
    out = pd.concat(frames, ignore_index=True)
    out[ts_id_col] = out[ts_id_col].astype(str)
    for col in out.columns:
        if col not in (ts_id_col, date_col, _SPLIT_COL) and pd.api.types.is_numeric_dtype(out[col]):
            out[col] = pd.to_numeric(out[col], errors="coerce").astype(float)
    return out

resolve_prior_tuning_artifact_uri

resolve_prior_tuning_artifact_uri(
    prior_run_id, model_type, *, settings, bq_client=None
)

Look up a previously saved stage_1_tuning_result_artifact_uri in BigQuery by run_id.

Checks forecast_metadata (model_artifact and best_params.stage_1_tuning_result_artifact_uri) and run_jobs.job_telemetry so users can warm-start a new Tabular Workflow run simply by setting "reuse_tuning_from_run_id": "<prior_run_id>" in model_params.

Source code in src/scale_forecasting/engines/automl_engine.py

def resolve_prior_tuning_artifact_uri(
    prior_run_id: str,
    model_type: str,
    *,
    settings: Settings,
    bq_client: Any = None,
) -> str | None:
    """Look up a previously saved ``stage_1_tuning_result_artifact_uri`` in BigQuery by ``run_id``.

    Checks ``forecast_metadata`` (``model_artifact`` and
    ``best_params.stage_1_tuning_result_artifact_uri``) and ``run_jobs.job_telemetry`` so users can
    warm-start a new Tabular Workflow run simply by setting
    ``"reuse_tuning_from_run_id": "<prior_run_id>"`` in ``model_params``.
    """
    from google.cloud import bigquery

    client = bq_client or bigquery_client(settings.project_id)
    meta_table = settings.registry_table_ref("forecast_metadata")
    sql = (
        f"SELECT model_artifact, TO_JSON_STRING(best_params) AS best_params_json "
        f"FROM `{meta_table}` "
        f"WHERE run_id = @run_id AND model_type = @model_type "
        f"LIMIT 1"
    )
    job_config = bigquery.QueryJobConfig(
        query_parameters=[
            bigquery.ScalarQueryParameter("run_id", "STRING", prior_run_id),
            bigquery.ScalarQueryParameter("model_type", "STRING", model_type),
        ]
    )
    try:
        rows = list(client.query(sql, job_config=job_config).result())
    except Exception as exc:  # noqa: BLE001
        _log.warning(
            "Could not query forecast_metadata for reuse_tuning_from_run_id=%r (%s)",
            prior_run_id,
            exc,
        )
        return None

    for row in rows:
        if row.get("best_params_json"):
            try:
                bp = json.loads(row["best_params_json"])
                uri = bp.get("stage_1_tuning_result_artifact_uri")
                if isinstance(uri, str) and uri.startswith("gs://"):
                    return uri
            except (ValueError, TypeError):
                pass
        art = row.get("model_artifact")
        if isinstance(art, str) and art.startswith("gs://") and "tuning" in art:
            return art
    return None

resolve_quantiles

resolve_quantiles(intervals, authored_quantiles=None)

Resolve the quantile list (max 5, including 0.5) for quantile objectives (pure).

Source code in src/scale_forecasting/engines/automl_engine.py

def resolve_quantiles(intervals: list[float], authored_quantiles: Any = None) -> list[float]:
    """Resolve the quantile list (max 5, including 0.5) for quantile objectives (pure)."""
    if isinstance(authored_quantiles, list) and authored_quantiles:
        qs = sorted({round(float(q), 4) for q in authored_quantiles} | {0.5})
        return qs[:5]
    q_set: set[float] = {0.5}
    for level in intervals:
        alpha = (1.0 - float(level)) / 2.0
        q_set.add(round(alpha, 4))
        q_set.add(round(1.0 - alpha, 4))
    return sorted(q_set)[:5]

run

run(
    cfg,
    *,
    models=None,
    manage_header=True,
    settings=None,
    job_id=None,
)

Run all selected automl models on Vertex AI and write results to BigQuery.

Source code in src/scale_forecasting/engines/automl_engine.py

def run(
    cfg: RunConfig,
    *,
    models: list[str] | None = None,
    manage_header: bool = True,
    settings: Settings | None = None,
    job_id: str | None = None,
) -> dict[str, Any]:
    """Run all selected ``automl`` models on Vertex AI and write results to BigQuery."""
    from ..registry.cells import write_cells
    from ..registry.lifecycle import run_header
    from ..settings import Settings as _Settings
    from .ray_engine import _assert_source_supports_folds, _read_source_series

    settings = settings or _Settings.resolve()
    run_id = make_run_id(cfg)
    executed_models = [
        m for m in (models if models is not None else cfg.models) if get_model(m).family == "automl"
    ]
    if not executed_models:
        return {"run_id": run_id, "n_succeeded": 0, "n_failed": 0, "models": {}}

    with run_header(cfg, run_id, settings=settings, manage=manage_header) as hdr:
        panel_df = _read_source_series(cfg, settings)
        _assert_source_supports_folds(panel_df, cfg)

        future_df: pd.DataFrame | None = None
        fut_table = getattr(cfg.data, "future_covariates_table", None)
        if fut_table:
            bq_client = bigquery_client(settings.project_id)
            fut_tbl = fut_table if "." in fut_table else settings.table_ref(fut_table)
            future_df = bq_client.query(f"SELECT * FROM `{fut_tbl}`").to_dataframe()

        all_cells: list[CellResult] = []
        model_telemetry: dict[str, Any] = {}
        for model_type in executed_models:
            cells, telem = execute_automl_model_cells(
                panel_df,
                model_type,
                cfg,
                run_id=run_id,
                settings=settings,
                future_df=future_df,
                job_id=job_id,
            )
            write_cells(cells, settings=settings)
            all_cells.extend(cells)
            model_telemetry[model_type] = telem

        n_ok = sum(1 for c in all_cells if c.status == "ok")
        n_fail = len(all_cells) - n_ok
        status = "COMPLETED" if n_fail == 0 else ("PARTIAL" if n_ok > 0 else "FAILED")
        n_series = int(panel_df[cfg.data.ts_id_col].nunique()) if not panel_df.empty else 0
        hdr.finalize(status=status, n_series=n_series)

    return {
        "run_id": run_id,
        "n_succeeded": n_ok,
        "n_failed": n_fail,
        "models": model_telemetry,
    }

training_job_class

training_job_class(aiplatform_module, model_type)

The trainer class TRAINING_JOB_CLASSES names for model_type, from the given SDK module.

Raises EngineError — never substitutes — when the SDK lacks the class: a user who asked for vertex_seq2seq must get seq2seq or an error, not L2L under a seq2seq label.

Source code in src/scale_forecasting/engines/automl_engine.py

def training_job_class(aiplatform_module: Any, model_type: str) -> Any:
    """The trainer class `TRAINING_JOB_CLASSES` names for ``model_type``, from the given SDK module.

    Raises `EngineError` — never substitutes — when the SDK lacks the class: a user who asked for
    ``vertex_seq2seq`` must get seq2seq or an error, not L2L under a seq2seq label.
    """
    base_model = base_model_name(model_type)
    try:
        class_name = TRAINING_JOB_CLASSES[base_model]
    except KeyError:
        raise EngineError(
            f"'{model_type}' has no training_job trainer; choose one of "
            f"{sorted(TRAINING_JOB_CLASSES)} or use automl_mode='tabular_workflow'."
        ) from None
    trainer = getattr(aiplatform_module, class_name, None)
    if trainer is None:
        raise EngineError(
            f"The installed google-cloud-aiplatform has no '{class_name}' (needed for "
            f"'{model_type}' in automl_mode='training_job'). Upgrade the SDK or use "
            "automl_mode='tabular_workflow'."
        )
    return trainer

training_job_kwargs

training_job_kwargs(plan, job_prefix)

Constructor kwargs for the aiplatform.*ForecastingTrainingJob of training_job mode.

Pure, and named once so a test can bind them against the installed SDK's signature. The keyword vocabulary here is the SDK's, not ours — column_specs maps each column to one of AutoMLTabularTrainingJob.column_data_types — and a mocked trainer accepts a misspelling the real class rejects at the first live run.

Source code in src/scale_forecasting/engines/automl_engine.py

def training_job_kwargs(plan: AutoMLModelPlan, job_prefix: str) -> dict[str, Any]:
    """Constructor kwargs for the ``aiplatform.*ForecastingTrainingJob`` of ``training_job`` mode.

    Pure, and named once so a test can bind them against the installed SDK's signature. The
    keyword vocabulary here is the SDK's, not ours — ``column_specs`` maps each column to one of
    ``AutoMLTabularTrainingJob.column_data_types`` — and a mocked trainer accepts a misspelling the
    real class rejects at the first live run.
    """
    return {
        "display_name": f"sf-train-{job_prefix}",
        "optimization_objective": plan.optimization_objective,
        "column_specs": {spec["column_name"]: spec["data_type"] for spec in plan.column_specs},
    }

training_run_kwargs

training_run_kwargs(plan, cfg, dataset)

run() kwargs for the training job training_job_kwargs constructs (pure).

Quantiles are only meaningful to the quantile-loss objective; the SDK rejects them otherwise, so they are passed as None for every other objective.

Source code in src/scale_forecasting/engines/automl_engine.py

def training_run_kwargs(plan: AutoMLModelPlan, cfg: RunConfig, dataset: Any) -> dict[str, Any]:
    """``run()`` kwargs for the training job `training_job_kwargs` constructs (pure).

    Quantiles are only meaningful to the quantile-loss objective; the SDK rejects them otherwise,
    so they are passed as ``None`` for every other objective.
    """
    return {
        "dataset": dataset,
        "target_column": cfg.data.target_col,
        "time_column": cfg.data.date_col,
        "time_series_identifier_column": cfg.data.ts_id_col,
        "unavailable_at_forecast_columns": plan.unavailable_at_forecast_columns,
        "available_at_forecast_columns": plan.available_at_forecast_columns,
        "time_series_attribute_columns": plan.time_series_attribute_columns,
        "forecast_horizon": plan.horizon,
        "context_window": plan.context_window,
        "data_granularity_unit": plan.data_granularity_unit,
        "data_granularity_count": plan.data_granularity_count,
        "predefined_split_column_name": _SPLIT_COL,
        "budget_milli_node_hours": plan.train_budget_milli_node_hours,
        "quantiles": (
            plan.quantiles if plan.optimization_objective == "minimize-quantile-loss" else None
        ),
        "sync": True,
    }