Vertex AI AutoML & Tabular Workflows Submitter

automl_submit

Submit a run (or automl family slice) to Vertex AI AutoML / Tabular Workflows.

Dispatches the four managed Vertex AI AutoML Forecasting model plugins (vertex_l2l, vertex_tide, vertex_tft, vertex_seq2seq) through engines.automl_engine. Supports both:

1. Tabular Workflow for Forecasting (automl_mode = "tabular_workflow", default): Compiles and submits a Vertex AI PipelineJob via google_cloud_pipeline_components, exposing per-stage hardware overrides (feature_transform_engine_machine_type, trainer_machine_type, trainer_replica_count, trainer_accelerator_type, trainer_accelerator_count, max_num_trials, max_parallel_trial_count) and Stage-1 tuning artifact warm-start (stage_1_tuning_result_artifact_uri or reuse_tuning_from_run_id).
2. Managed AutoML TrainingJob (automl_mode = "training_job"): Submits a fully managed AutoMLForecastingTrainingJob / TimeSeriesDenseEncoderForecastingTrainingJob / TemporalFusionTransformerForecastingTrainingJob / SequenceToSequencePlusForecastingTrainingJob via the google-cloud-aiplatform SDK.

After training completes, executes a BatchPredictionJob with generate_explanation=True and writes calibrated forecasts, backtest OOF evaluations, Tier 1 series-level feature attributions (forecast_metadata.fit_diagnostics["feature_attributions"]), Tier 2 per-horizon-step local attributions (forecast_predictions.explanations), and GCS tuning/model artifact URIs (forecast_metadata.model_artifact and best_params) back to the BigQuery registry.

VertexAutoMLJobPlan dataclass

VertexAutoMLJobPlan(
    display_name,
    automl_mode,
    models,
    hardware,
    gpu_type,
    machine_type,
    max_workers,
    accelerator_count,
    config_uri="",
    service_account="",
)

Resolved execution summary for a vertex_automl family job (pure).

to_dict

to_dict()

JSON-safe summary for run_jobs.job_telemetry.$.vertex_automl.

Source code in src/scale_forecasting/automl_submit.py

def to_dict(self) -> dict[str, Any]:
    """JSON-safe summary for ``run_jobs.job_telemetry.$.vertex_automl``."""
    return {
        "display_name": self.display_name,
        "automl_mode": self.automl_mode,
        "models": list(self.models),
        "hardware": self.hardware,
        "gpu_type": self.gpu_type,
        "machine_type": self.machine_type,
        "max_workers": self.max_workers,
        "accelerator_count": self.accelerator_count,
        "config_uri": self.config_uri,
        "service_account": self.service_account,
    }

plan_automl_job

plan_automl_job(
    cfg,
    models=None,
    *,
    run_id=None,
    job_id=None,
    automl_mode=None,
    hardware=None,
    gpu_type=None,
    machine_type=None,
    max_workers=None,
    accelerator_count=None,
    config_uri="",
    service_account="",
)

Resolve a VertexAutoMLJobPlan from cfg and per-family overrides (pure).

Source code in src/scale_forecasting/automl_submit.py

def plan_automl_job(
    cfg: RunConfig,
    models: Sequence[str] | None = None,
    *,
    run_id: str | None = None,
    job_id: str | None = None,
    automl_mode: str | None = None,
    hardware: str | None = None,
    gpu_type: str | None = None,
    machine_type: str | None = None,
    max_workers: int | None = None,
    accelerator_count: int | None = None,
    config_uri: str = "",
    service_account: str = "",
) -> VertexAutoMLJobPlan:
    """Resolve a `VertexAutoMLJobPlan` from `cfg` and per-family overrides (pure)."""
    rid = run_id or make_run_id(cfg)
    selected = [
        m for m in (models if models is not None else cfg.models) if get_model(m).family == "automl"
    ]
    fc = cfg.resolve_family_compute("automl")
    eff_mode = automl_mode or fc.automl_mode or cfg.compute.automl_mode
    eff_hw = hardware or fc.hardware or "cpu"
    eff_gpu = (gpu_type or fc.gpu_type) if eff_hw == "gpu" else None
    eff_machine = (
        machine_type or fc.machine_type or ("g2-standard-8" if eff_hw == "gpu" else "n1-standard-8")
    )
    eff_workers = max_workers or fc.max_workers or fc.workers or 10
    eff_accel = (
        (accelerator_count if accelerator_count is not None else fc.accelerator_count)
        if eff_hw == "gpu"
        else 0
    )
    display_name = (
        vertex_automl_job_id(job_id) if job_id else vertex_automl_job_id(f"sf-{rid}-automl")
    )
    return VertexAutoMLJobPlan(
        display_name=display_name,
        automl_mode=eff_mode,
        models=tuple(selected),
        hardware=eff_hw,
        gpu_type=eff_gpu,
        machine_type=eff_machine,
        max_workers=int(eff_workers),
        accelerator_count=int(eff_accel),
        config_uri=config_uri,
        service_account=service_account,
    )

submit_automl

submit_automl(
    cfg,
    settings=None,
    infra=None,
    *,
    wait=True,
    models=None,
    job_id=None,
    manage_header=True,
    automl_mode=None,
    hardware=None,
    gpu_type=None,
    machine_type=None,
    max_workers=None,
    accelerator_count=None,
)

Execute automl models on Vertex AI and return (run_id, native_id, probe_handle).

Source code in src/scale_forecasting/automl_submit.py

def submit_automl(
    cfg: RunConfig,
    settings: Settings | None = None,
    infra: BatchInfra | None = None,
    *,
    wait: bool = True,
    models: list[str] | None = None,
    job_id: str | None = None,
    manage_header: bool = True,
    automl_mode: str | None = None,
    hardware: str | None = None,
    gpu_type: str | None = None,
    machine_type: str | None = None,
    max_workers: int | None = None,
    accelerator_count: int | None = None,
) -> tuple[str, str, ProbeHandle]:
    """Execute ``automl`` models on Vertex AI and return ``(run_id, native_id, probe_handle)``."""
    settings = settings or Settings.resolve()
    if infra is None:
        with contextlib.suppress(Exception):
            infra = BatchInfra.resolve()

    run_id = make_run_id(cfg)
    automl_models = [
        m for m in (models if models is not None else cfg.models) if get_model(m).family == "automl"
    ]
    if not automl_models:
        raise ConfigError(
            "submit_automl called with a config that has no 'automl' family models to execute."
        )

    code_bucket = ""
    service_account = ""
    if infra is not None:
        code_bucket = infra.code_bucket
        service_account = infra.compute_sa
    elif settings.warehouse_uri.startswith("gs://"):
        code_bucket = settings.warehouse_uri.removeprefix("gs://").split("/", 1)[0]

    config_uri = ""
    if code_bucket:
        with contextlib.suppress(Exception):
            config_uri = staging.stage_config(cfg, run_id, code_bucket)

    plan = plan_automl_job(
        cfg,
        automl_models,
        run_id=run_id,
        job_id=job_id,
        automl_mode=automl_mode,
        hardware=hardware,
        gpu_type=gpu_type,
        machine_type=machine_type,
        max_workers=max_workers,
        accelerator_count=accelerator_count,
        config_uri=config_uri,
        service_account=service_account,
    )

    profile = None
    with contextlib.suppress(Exception):
        profile = profile_for_run(cfg, settings=settings)

    sizing_entry = {
        "family": "automl",
        "runtime": "vertex_automl",
        "automl_mode": plan.automl_mode,
        "hardware": plan.hardware,
        "gpu_type": plan.gpu_type,
        "machine_type": plan.machine_type,
        "max_workers": plan.max_workers,
        "models": list(plan.models),
        "profile_source": getattr(profile, "source", "static") if profile is not None else "static",
    }
    with contextlib.suppress(Exception):
        merge_header_telemetry(
            run_id,
            {
                "sizing": {"automl": sizing_entry},
                "sizing_executed": {"automl": sizing_entry},
            },
            settings=settings,
        )

    handle = ProbeHandle(
        runtime="vertex_automl",
        native_id=plan.display_name,
        region=settings.region,
    )
    if job_id is not None:
        with contextlib.suppress(Exception):
            update_job(
                job_id,
                settings=settings,
                merge_telemetry={
                    "probe_handle": handle.to_blob(),
                    "vertex_automl": plan.to_dict(),
                },
            )

    import time

    _ = wait
    t_start = time.monotonic()
    outcome = automl_engine.run(
        cfg,
        models=automl_models,
        manage_header=manage_header,
        settings=settings,
        job_id=plan.display_name,
    )
    wall_s = round(max(0.0, time.monotonic() - t_start), 1)

    # If the pipeline produced a concrete resource_name or pipeline_job_id, update the handle.
    models_telem = outcome.get("models") or {}
    first_telem: dict[str, Any] = (
        next(iter(models_telem.values()), {}) if isinstance(models_telem, dict) else {}
    )
    resource_name = first_telem.get("pipeline_resource_name") or first_telem.get(
        "vertex_model_resource_name"
    )
    native_id = first_telem.get("pipeline_job_id") or plan.display_name
    final_handle = ProbeHandle(
        runtime="vertex_automl",
        native_id=str(native_id),
        region=settings.region,
        resource_name=str(resource_name) if resource_name else None,
    )
    prov_s = first_telem.get("provision_s")
    exec_s = first_telem.get("execute_s")
    if exec_s is None:
        exec_s = wall_s
    if prov_s is None:
        prov_s = round(max(0.0, wall_s - float(exec_s)), 1)
    if job_id is not None:
        with contextlib.suppress(Exception):
            update_job(
                job_id,
                settings=settings,
                merge_telemetry={
                    "probe_handle": final_handle.to_blob(),
                    "platform": {
                        "provision_s": prov_s,
                        "execute_s": exec_s,
                        "total_wall_s": wall_s,
                    },
                    "vertex_automl": {
                        **plan.to_dict(),
                        "models_telemetry": models_telem,
                    },
                },
            )

    return run_id, final_handle.native_id, final_handle