diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
new file mode 100644
index 0000000..b06157a
--- /dev/null
+++ b/.github/workflows/ci.yml
@@ -0,0 +1,38 @@
+name: Dataproc Custom Image
+
+on:
+  push:
+    branches: [ main ]
+  pull_request:
+    branches: [ main ]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.10", "3.11", "3.12"]
+
+    steps:
+    - name: Checkout Repository
+      uses: actions/checkout@v6
+
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v6
+      with:
+        python-version: ${{ matrix.python-version }}
+        cache: 'pip'
+
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install -r requirements.txt
+
+    - name: Run Ruff check (Linter)
+      run: ruff check .
+
+    - name: Run Ruff format check
+      run: ruff format --check .
+
+    - name: Run Unit Tests
+      run: python -m unittest discover -s tests -p "test_*.py"
diff --git a/README.md b/README.md
index e13a728..3d6a263 100644
--- a/README.md
+++ b/README.md
@@ -133,6 +133,10 @@ python generate_custom_image.py \
     default value of 300 seconds will be used.
 *   **--dry-run**: Dry run mode which only validates input and generates
     workflow script without creating image. Disabled by default.
+*   **--execution-engine**: The orchestration backend used to manage the custom
+    image creation. Supported values are `cli` (runs legacy generated shell-script
+    via local gcloud commands) and `api` (uses native Google Cloud Python clients).
+    Defaults to `cli`.
 *   **--trusted-cert**: a certificate in DER format to be inserted
     into the custom image's EFI boot sector.  Can be generated by
     reading examples/secure-boot/README.md.  This argument is mutually
@@ -183,13 +187,13 @@ cluster property values in the file.
 
 #### Create a custom image
 
-Create a custom image with name `custom-image-1-5-9` with Dataproc version
-`1.5.9-debian10`:
+Create a custom image with name `custom-image-2-2-0` with Dataproc version
+`2.2.0-debian12`:
 
 ```shell
 python generate_custom_image.py \
-    --image-name custom-image-1-5-9 \
-    --dataproc-version 1.5.9-debian10 \
+    --image-name custom-image-2-2-0 \
+    --dataproc-version 2.2.0-debian12 \
     --customization-script ~/custom-script.sh \
     --metadata 'key1=value1,key2=value2' \
     --zone us-central1-f \
@@ -200,10 +204,37 @@ python generate_custom_image.py \
 
 ```shell
 python generate_custom_image.py \
-    --image-name custom-image-1-5-9 \
-    --dataproc-version 1.5.9-debian10 \
+    --image-name custom-image-2-2-0 \
+    --dataproc-version 2.2.0-debian12 \
     --customization-script ~/custom-script.sh \
     --zone us-central1-f \
     --gcs-bucket gs://my-test-bucket \
     --no-smoke-test
 ```
+
+### API-based Execution Engine (Optional)
+
+The tool supports an optional API-based execution engine that communicates directly with Google Cloud services using official Python Client Libraries instead of generating and running local shell scripts.
+
+#### Setup
+
+To use the API-based execution engine, install the required packages:
+
+```shell
+pip install -r requirements.txt
+```
+
+#### Running in API Mode
+
+To invoke the custom image generator using the API engine, pass `--execution-engine api`:
+
+```shell
+python generate_custom_image.py \
+    --image-name custom-image-2-2-0 \
+    --dataproc-version 2.2.0-debian12 \
+    --customization-script ~/custom-script.sh \
+    --zone us-central1-f \
+    --gcs-bucket gs://my-test-bucket \
+    --execution-engine api
+```
+
diff --git a/custom_image_utils/api_execution_engine.py b/custom_image_utils/api_execution_engine.py
new file mode 100644
index 0000000..2f93ce1
--- /dev/null
+++ b/custom_image_utils/api_execution_engine.py
@@ -0,0 +1,916 @@
+# Copyright 2026 Google LLC and contributors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#            http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""API-based execution engine implementation using Google Cloud Python client libraries."""
+
+import datetime
+from dataclasses import dataclass
+import logging
+import os
+import re
+import sys
+import time
+
+import google.auth
+from google.api_core.exceptions import (
+    NotFound,
+    GoogleAPIError,
+    PermissionDenied,
+    Conflict,
+    DeadlineExceeded,
+)
+from google.api_core.retry import Retry
+from google.cloud import compute_v1
+from google.cloud import storage
+
+from custom_image_utils.execution_engine import ExecutionEngine
+from custom_image_utils.compute_operation_helper import ComputeOperationHelper
+
+
+def _is_transient(e):
+    from google.api_core.exceptions import BadRequest
+
+    if isinstance(e, (NotFound, PermissionDenied, Conflict, BadRequest)):
+        return False
+    return isinstance(e, (GoogleAPIError, DeadlineExceeded))
+
+
+_DEFAULT_RETRY = Retry(
+    initial=1.0,
+    maximum=10.0,
+    multiplier=2.0,
+    predicate=_is_transient,
+)
+
+
+def _parse_dataproc_version(version_str):
+    """Parses dataproc version string like '2.1.115-debian11' or '2-1-115-debian11' to integer tuple (2, 1, 115) for sorting."""
+    if not version_str:
+        return 0, 0, 0
+    # Normalize hyphens to dots to handle both formats
+    normalized = version_str.replace("-", ".")
+    parts = []
+    for p in normalized.split("."):
+        try:
+            parts.append(int(p))
+        except ValueError:
+            break
+    while len(parts) < 3:
+        parts.append(0)
+    return tuple(parts[:3])
+
+
+@dataclass
+class BuildState:
+    """State of the VM and disk provisioning/creation workflow."""
+
+    disk_created: bool = False
+    vm_created: bool = False
+    build_succeeded: bool = False
+    build_failed: bool = False
+
+
+_LOG = logging.getLogger(__name__)
+_LOG.setLevel(logging.INFO)
+
+_IMAGE_PATH = "projects/{}/global/images/{}"
+_IMAGE_URI = re.compile(
+    r"^(https://www\.googleapis\.com/compute/([^/]+)/)?projects/([^/]+)/global/images/([^/]+)$"
+)
+_IMAGE_FAMILY_PATH = "projects/{}/global/images/family/{}"
+_IMAGE_FAMILY_URI = re.compile(
+    r"^(https://www\.googleapis\.com/compute/([^/]+)/)?projects/([^/]+)/global/images/family/([^/]+)$"
+)
+
+
+def _has_build_signal(contents: str, signal: str) -> bool:
+    """Checks if contents has the build signal in a non-trace and non-quoted way."""
+    target = f"startup-script: {signal}"
+    for line in contents.splitlines():
+        if target in line:
+            if "+ " in line:
+                continue
+            if f'"{target}' in line or f"'{target}" in line or f'\\"{target}' in line:
+                continue
+            return True
+    return False
+
+
+class ApiExecutionEngine(ExecutionEngine):
+    """Execution engine that uses Google Cloud Python client libraries."""
+
+    def __init__(self, credentials=None):
+        self.credentials = credentials
+        self.compute_helper = ComputeOperationHelper(credentials=credentials)
+        self.images_client = compute_v1.ImagesClient(credentials=credentials)
+        self.disks_client = compute_v1.DisksClient(credentials=credentials)
+        self.instances_client = compute_v1.InstancesClient(credentials=credentials)
+        self.storage_client = storage.Client(credentials=credentials)
+
+    def _get_default_project(self):
+        """Gets default project ID from authenticated credentials."""
+        if self.credentials and getattr(self.credentials, "project_id", None):
+            return self.credentials.project_id
+        _, project_id = google.auth.default()
+        if not project_id:
+            raise RuntimeError(
+                "Cannot find default Google Cloud project ID. "
+                "Please verify your credentials or set --project-id."
+            )
+        return project_id
+
+    def infer_args(self, args):
+        """Infers missing command-line arguments using GCP client APIs."""
+        _LOG.info("Inferring arguments using API Engine...")
+
+        # 1. Project ID
+        if not args.project_id:
+            args.project_id = self._get_default_project()
+
+        # Validate and format Zone & Region (Point 12)
+        if not args.zone:
+            raise RuntimeError("Zone must be specified.")
+        zone_pattern = re.compile(r"^[a-z0-9-]+-[a-z]$")
+        if not zone_pattern.match(args.zone):
+            raise RuntimeError(
+                f"Invalid zone format: {args.zone}. Expected format like us-central1-a"
+            )
+
+        region_from_zone = "-".join(args.zone.split("-")[:-1])
+        region_pattern = re.compile(r"^[a-z0-9-]+$")
+        if not region_pattern.match(region_from_zone):
+            raise RuntimeError(f"Extracted region is invalid: {region_from_zone}")
+
+        # 2. Base Image
+        if args.base_image_uri:
+            m = _IMAGE_URI.match(args.base_image_uri)
+            project, image_name = m.group(3), m.group(4)
+            args.dataproc_base_image = _IMAGE_PATH.format(project, image_name)
+
+            # Describe base image to get dataproc version label
+            img = self.images_client.get(
+                project=project, image=image_name, retry=_DEFAULT_RETRY
+            )
+            args.dataproc_version = (img.labels or {}).get("goog-dataproc-version", "")
+
+        elif args.dataproc_version:
+            # Find base image path by dataproc version
+            parsed_version = args.dataproc_version.split(".")
+            major_version = parsed_version[0]
+            if len(parsed_version) == 2:
+                # e.g., 1.5-debian10 -> query READY images and filter in Python
+                minor_version = parsed_version[1].split("-")[0]
+                version_filter = parsed_version[1].replace("-", r"-\d+-", 1)
+                label_regex = re.compile(f"^{parsed_version[0]}-{version_filter}$")
+                filter_expr = 'status = "READY"'
+            else:
+                major_version = parsed_version[0]
+                minor_version = parsed_version[1]
+                version_str = (
+                    f"{parsed_version[0]}-{parsed_version[1]}-{parsed_version[2]}"
+                )
+                label_regex = None
+                filter_expr = f'labels.goog-dataproc-version = "{version_str}" AND status = "READY"'
+
+            # List matching Dataproc base images
+            images = list(
+                self.images_client.list(
+                    request={"project": "cloud-dataproc", "filter": filter_expr},
+                    retry=_DEFAULT_RETRY,
+                )
+            )
+
+            # Sort by parsed version descending, then by creationTimestamp descending (Point 13)
+            images.sort(
+                key=lambda x: (
+                    _parse_dataproc_version(x.labels.get("goog-dataproc-version", "")),
+                    x.creation_timestamp or "",
+                ),
+                reverse=True,
+            )
+
+            expected_prefix = f"dataproc-{major_version}-{minor_version}"
+            all_images_for_version = {}
+            image_versions = []
+
+            for img in images:
+                # Local Python filtering to avoid fragile/non-standard API filters
+                if "-eap" in img.name:
+                    continue
+                if not img.name.startswith(expected_prefix):
+                    continue
+                ver = img.labels.get("goog-dataproc-version")
+                if not ver:
+                    continue
+                if label_regex and not label_regex.match(ver):
+                    continue
+
+                if ver not in all_images_for_version:
+                    all_images_for_version[ver] = [
+                        _IMAGE_PATH.format("cloud-dataproc", img.name)
+                    ]
+                    image_versions.append(ver)
+                else:
+                    all_images_for_version[ver].append(
+                        _IMAGE_PATH.format("cloud-dataproc", img.name)
+                    )
+
+            if not image_versions:
+                raise RuntimeError(
+                    f"Cannot find dataproc base image with version {args.dataproc_version}"
+                )
+
+            latest_ver = image_versions[0]
+            if len(all_images_for_version[latest_ver]) > 1:
+                raise RuntimeError(
+                    "Found more than one image for latest dataproc version."
+                    f" Images: {all_images_for_version[latest_ver]}"
+                )
+
+            args.dataproc_base_image = all_images_for_version[latest_ver][0]
+            args.dataproc_version = latest_ver
+
+        elif args.base_image_family:
+            m = _IMAGE_FAMILY_URI.match(args.base_image_family)
+            project, family_name = m.group(3), m.group(4)
+            args.dataproc_base_image = _IMAGE_FAMILY_PATH.format(project, family_name)
+
+            # Describe latest image from the family
+            img = self.images_client.get_from_family(
+                project=project, family=family_name, retry=_DEFAULT_RETRY
+            )
+            args.dataproc_version = img.labels.get("goog-dataproc-version", "")
+        else:
+            raise RuntimeError(
+                "Neither --dataproc-version nor --base-image-uri nor"
+                " --base-image-family is specified."
+            )
+
+        # 3. OAuth config (None or formatted string)
+        if args.oauth:
+            args.oauth_path = os.path.abspath(args.oauth)
+        else:
+            args.oauth_path = None
+
+        # 4. Network and subnetwork configuration (Point 3)
+        if not args.network and not args.subnetwork:
+            args.network = f"projects/{args.project_id}/global/networks/default"
+
+        # Expand Network Uri
+        if args.network and not args.network.startswith("projects/"):
+            if args.network.startswith("global/networks/"):
+                args.network = f"projects/{args.project_id}/{args.network}"
+            elif "/" not in args.network:
+                args.network = (
+                    f"projects/{args.project_id}/global/networks/{args.network}"
+                )
+
+        # Expand Subnetwork Uri
+        if args.subnetwork and not args.subnetwork.startswith("projects/"):
+            if args.subnetwork.startswith("regions/"):
+                args.subnetwork = f"projects/{args.project_id}/{args.subnetwork}"
+            elif "/" not in args.subnetwork:
+                args.subnetwork = f"projects/{args.project_id}/regions/{region_from_zone}/subnetworks/{args.subnetwork}"
+
+        args.shutdown_timer_in_sec = args.shutdown_instance_timer_sec
+
+        _LOG.info("Returned Dataproc base image: %s", args.dataproc_base_image)
+        _LOG.info("Returned Dataproc version   : %s", args.dataproc_version)
+        print(f"Returned Dataproc base image: {args.dataproc_base_image}")
+        print(f"Returned Dataproc version   : {args.dataproc_version}")
+
+    def perform_sanity_checks(self, args):
+        """Checks if the target image already exists using Images API client."""
+        _LOG.info("Performing sanity checks using API Client...")
+        try:
+            self.images_client.get(
+                project=args.project_id, image=args.image_name, retry=_DEFAULT_RETRY
+            )
+            raise RuntimeError(f"Image {args.image_name} already exists.")
+        except NotFound:
+            # Image does not exist, which is expected.
+            pass
+        except GoogleAPIError as e:
+            raise RuntimeError(f"Error describing image {args.image_name}: {e}")
+        _LOG.info("Passed sanity checks...")
+
+    def create_image(self, args):
+        """Executes the custom image creation workflow using API Clients."""
+        if args.dry_run:
+            _LOG.info("Dry-run mode: Skipping image creation.")
+            return
+
+        # Initialize runtime identifiers
+        if "run_id" not in vars(args):
+            args.run_id = "custom-image-{image_name}-{timestamp}".format(
+                image_name=args.image_name,
+                timestamp=datetime.datetime.now().strftime("%Y%m%d-%H%M%S"),
+            )
+        gcs_bucket_clean = args.gcs_bucket.replace("gs://", "").strip("/")
+        if "/" in gcs_bucket_clean:
+            args.bucket_name, prefix_path = gcs_bucket_clean.split("/", 1)
+            prefix_path = prefix_path.strip("/")
+            args.gcs_base_path = f"{prefix_path}/{args.run_id}"
+        else:
+            args.bucket_name = gcs_bucket_clean
+            args.gcs_base_path = args.run_id
+
+        args.custom_sources_path = (
+            f"gs://{args.bucket_name}/{args.gcs_base_path}/sources"
+        )
+        args.log_dir = f"/tmp/{args.run_id}/logs"
+        args.gcs_log_dir = f"gs://{args.bucket_name}/{args.gcs_base_path}/logs"
+
+        os.makedirs(args.log_dir, exist_ok=True)
+        local_log_file = os.path.join(args.log_dir, "startup-script.log")
+
+        # Upload customizing sources to GCS
+        _LOG.info("Uploading files to GCS bucket...")
+        all_sources = {
+            "run.sh": "startup_script/run.sh",
+            "gce-proxy-setup.sh": "startup_script/gce-proxy-setup.sh",
+        }
+        # Upload any non-GCS extra sources
+        for target_name, path in args.extra_sources.items():
+            if not path.startswith("gs://"):
+                all_sources[target_name] = path
+
+        bucket = self.storage_client.bucket(args.bucket_name)
+        for target_name, local_path in all_sources.items():
+            blob_path = f"{args.gcs_base_path}/sources/{target_name}"
+            blob = bucket.blob(blob_path)
+            blob.upload_from_filename(local_path, retry=_DEFAULT_RETRY)
+            _LOG.info("Uploaded %s to %s", local_path, blob_path)
+
+            # Post-upload verification check (Point 9)
+            if not blob.exists(retry=_DEFAULT_RETRY):
+                raise RuntimeError(
+                    f"Failed to verify upload of {local_path} to GCS (blob does not exist)."
+                )
+
+        # Handle customization script (local vs. gs://) (Point 10)
+        init_actions_blob_path = f"{args.gcs_base_path}/sources/init_actions.sh"
+        if args.customization_script.startswith("gs://"):
+            # Cloud-to-Cloud copy! No local download needed.
+            _LOG.info(
+                "Copying remote GCS customization script: %s", args.customization_script
+            )
+            src_bucket_name, src_blob_name = args.customization_script.replace(
+                "gs://", ""
+            ).split("/", 1)
+            src_bucket = self.storage_client.bucket(src_bucket_name)
+            src_blob = src_bucket.blob(src_blob_name)
+            try:
+                src_bucket.copy_blob(
+                    src_blob, bucket, init_actions_blob_path, retry=_DEFAULT_RETRY
+                )
+            except (
+                GoogleAPIError,
+                NotFound,
+                PermissionDenied,
+                Conflict,
+                DeadlineExceeded,
+            ) as e:
+                raise RuntimeError(
+                    f"Error copying customization script {args.customization_script} to {init_actions_blob_path}: {e}. "
+                    "Please ensure your service account has storage.objects.get permission on the source bucket "
+                    "and storage.objects.create permission on the target bucket."
+                ) from e
+            _LOG.info(
+                "Copied remote GCS customization script to %s", init_actions_blob_path
+            )
+
+            # Verify GCS copy succeeded
+            dest_blob = bucket.blob(init_actions_blob_path)
+            if not dest_blob.exists(retry=_DEFAULT_RETRY):
+                raise RuntimeError(
+                    f"Failed to verify GCS copy of customization script to {init_actions_blob_path}."
+                )
+        else:
+            # Upload local file
+            blob = bucket.blob(init_actions_blob_path)
+            blob.upload_from_filename(args.customization_script, retry=_DEFAULT_RETRY)
+            _LOG.info(
+                "Uploaded local customization script %s to %s",
+                args.customization_script,
+                init_actions_blob_path,
+            )
+
+            # Verify local upload succeeded
+            if not blob.exists(retry=_DEFAULT_RETRY):
+                raise RuntimeError(
+                    f"Failed to verify upload of customization script {args.customization_script} to GCS."
+                )
+
+        # Handle GCS extra sources if any
+        for target_name, path in args.extra_sources.items():
+            if path.startswith("gs://"):
+                _LOG.info("Copying remote GCS extra source: %s", path)
+                src_bucket_name, src_blob_name = path.replace("gs://", "").split("/", 1)
+                src_bucket = self.storage_client.bucket(src_bucket_name)
+                src_blob = src_bucket.blob(src_blob_name)
+                try:
+                    src_bucket.copy_blob(
+                        src_blob,
+                        bucket,
+                        f"{args.gcs_base_path}/sources/{target_name}",
+                        retry=_DEFAULT_RETRY,
+                    )
+                except (
+                    GoogleAPIError,
+                    NotFound,
+                    PermissionDenied,
+                    Conflict,
+                    DeadlineExceeded,
+                ) as e:
+                    raise RuntimeError(
+                        f"Error copying GCS extra source {path} to sources/{target_name}: {e}. "
+                        "Please ensure your service account has storage.objects.get permission on the source bucket "
+                        "and storage.objects.create permission on the target bucket."
+                    ) from e
+                _LOG.info(
+                    "Copied remote GCS extra source %s to sources/%s", path, target_name
+                )
+
+                # Verify copy succeeded
+                extra_blob = bucket.blob(f"{args.gcs_base_path}/sources/{target_name}")
+                if not extra_blob.exists(retry=_DEFAULT_RETRY):
+                    raise RuntimeError(
+                        f"Failed to verify GCS copy of extra source {path} to GCS."
+                    )
+
+        # Resolve zone and region
+        region = "-".join(args.zone.split("-")[:-1])
+        disk_name = f"{args.image_name}-install"
+        instance_name = f"{args.image_name}-install"
+
+        state = BuildState()
+
+        try:
+            # Create Compute Disk from base image
+            self._create_disk(args, disk_name)
+            state.disk_created = True
+
+            # Create VM Instance
+            self._create_vm(args, instance_name, disk_name, region)
+            state.vm_created = True
+
+            # Monitor Serial Logs
+            self._monitor_build(args, instance_name, local_log_file, state)
+
+            # Check outcome status
+            if not state.build_succeeded:
+                raise RuntimeError(
+                    "Custom image build failed. See logs at {} or GCS {}.".format(
+                        local_log_file, args.gcs_log_dir
+                    )
+                )
+
+            # Ensure the VM is fully stopped before creating the image from its disk
+            inst = self.instances_client.get(
+                project=args.project_id,
+                zone=args.zone,
+                instance=instance_name,
+                retry=_DEFAULT_RETRY,
+            )
+            if inst.status not in ("TERMINATED", "STOPPED"):
+                _LOG.info(
+                    "Stopping VM instance %s to release disk for imaging...",
+                    instance_name,
+                )
+                op = self.instances_client.stop(
+                    project=args.project_id,
+                    zone=args.zone,
+                    instance=instance_name,
+                    retry=_DEFAULT_RETRY,
+                )
+                self.compute_helper.wait_for_zone_operation(
+                    args.project_id, args.zone, op.name
+                )
+
+            # Create Final custom image from the disk
+            self._create_final_image(args, disk_name)
+
+            # Shutdown and delete the VM instance
+            _LOG.info("Deleting VM instance %s...", instance_name)
+            op = self.instances_client.delete(
+                project=args.project_id,
+                zone=args.zone,
+                instance=instance_name,
+                retry=_DEFAULT_RETRY,
+            )
+            self.compute_helper.wait_for_zone_operation(
+                args.project_id, args.zone, op.name
+            )
+            state.vm_created = False
+
+        finally:
+            self._cleanup(args, instance_name, disk_name, state)
+            self._upload_logs(args, bucket)
+
+    def _create_disk(self, args, disk_name):
+        """Creates boot disk from Dataproc base image."""
+        _LOG.info(
+            "Creating boot disk %s from base image %s...",
+            disk_name,
+            args.dataproc_base_image,
+        )
+        disk_body = compute_v1.Disk(
+            name=disk_name,
+            source_image=args.dataproc_base_image,
+            type_=f"zones/{args.zone}/diskTypes/pd-ssd",
+            size_gb=args.disk_size,
+        )
+        try:
+            op = self.disks_client.insert(
+                project=args.project_id,
+                zone=args.zone,
+                disk_resource=disk_body,
+                retry=_DEFAULT_RETRY,
+            )
+            self.compute_helper.wait_for_zone_operation(
+                args.project_id, args.zone, op.name
+            )
+        except (
+            GoogleAPIError,
+            NotFound,
+            PermissionDenied,
+            Conflict,
+            DeadlineExceeded,
+        ) as e:
+            raise RuntimeError(f"Error creating boot disk {disk_name}: {e}")
+
+    def _create_vm(self, args, instance_name, disk_name, region):
+        """Creates the VM instance that runs customization scripts."""
+        _LOG.info(
+            "Creating VM instance %s to run customization script...", instance_name
+        )
+
+        # Build metadata items
+        metadata_items = [
+            compute_v1.Items(
+                key="shutdown-timer-in-sec", value=str(args.shutdown_timer_in_sec)
+            ),
+            compute_v1.Items(key="custom-sources-path", value=args.custom_sources_path),
+            compute_v1.Items(key="universe-domain", value=args.universe_domain),
+            compute_v1.Items(key="dataproc-region", value=region),
+        ]
+        if args.dataproc_version:
+            metadata_items.append(
+                compute_v1.Items(
+                    key="dataproc_dataproc_version", value=args.dataproc_version
+                )
+            )
+        # Process user customization metadata (Point 4)
+        if args.metadata:
+            # Match key=value where value can be double-quoted or unquoted
+            for match in re.finditer(
+                r'([^=,\s]+)=(?:"([^"]*)"|([^,]+))', args.metadata
+            ):
+                k = match.group(1)
+                v = match.group(2) if match.group(2) is not None else match.group(3)
+                metadata_items.append(compute_v1.Items(key=k, value=v))
+
+        # Use startup-script-url pointing to GCS to avoid ~256KB metadata limits (Point 11)
+        startup_script_url = (
+            f"gs://{args.bucket_name}/{args.gcs_base_path}/sources/run.sh"
+        )
+        metadata_items.append(
+            compute_v1.Items(key="startup-script-url", value=startup_script_url)
+        )
+
+        # Build network interface
+        network_interface = compute_v1.NetworkInterface()
+        if args.subnetwork:
+            network_interface.subnetwork = args.subnetwork
+        else:
+            network_interface.network = args.network
+        if not args.no_external_ip:
+            # Add access config to assign external IP address
+            network_interface.access_configs = [
+                compute_v1.AccessConfig(
+                    name="External NAT",
+                    type_="ONE_TO_ONE_NAT",
+                )
+            ]
+
+        # Build Boot Disk attachment
+        boot_disk = compute_v1.AttachedDisk(
+            auto_delete=True,
+            boot=True,
+            mode="READ_WRITE",
+            source=f"projects/{args.project_id}/zones/{args.zone}/disks/{disk_name}",
+        )
+
+        # Build instance specs
+        instance_resource = compute_v1.Instance(
+            name=instance_name,
+            machine_type=f"zones/{args.zone}/machineTypes/{args.machine_type}",
+            disks=[boot_disk],
+            network_interfaces=[network_interface],
+            metadata=compute_v1.Metadata(items=metadata_items),
+            service_accounts=[
+                compute_v1.ServiceAccount(
+                    email=args.service_account,
+                    scopes=["https://www.googleapis.com/auth/cloud-platform"],
+                )
+            ],
+        )
+
+        # Add Accelerator configurations
+        if args.accelerator:
+            # type=nvidia-tesla-v100,count=2
+            acc_type = "nvidia-tesla-v100"
+            acc_count = 1
+            for item in args.accelerator.split(","):
+                if "=" in item:
+                    k, v = item.split("=", 1)
+                    if k == "type":
+                        acc_type = v
+                    elif k == "count":
+                        try:
+                            acc_count = int(v)
+                        except ValueError:
+                            raise RuntimeError(
+                                f"Invalid accelerator count: {v}. Must be an integer."
+                            )
+            instance_resource.guest_accelerators = [
+                compute_v1.AcceleratorConfig(
+                    accelerator_type=f"zones/{args.zone}/acceleratorTypes/{acc_type}",
+                    accelerator_count=acc_count,
+                )
+            ]
+            # Required scheduling for GPU instances
+            instance_resource.scheduling = compute_v1.Scheduling(
+                on_host_maintenance="TERMINATE"
+            )
+
+        try:
+            op = self.instances_client.insert(
+                project=args.project_id,
+                zone=args.zone,
+                instance_resource=instance_resource,
+                retry=_DEFAULT_RETRY,
+            )
+            self.compute_helper.wait_for_zone_operation(
+                args.project_id, args.zone, op.name
+            )
+        except (
+            GoogleAPIError,
+            NotFound,
+            PermissionDenied,
+            Conflict,
+            DeadlineExceeded,
+        ) as e:
+            raise RuntimeError(f"Error creating VM instance {instance_name}: {e}")
+
+    def _monitor_build(self, args, instance_name, local_log_file, state):
+        """Monitors the customization build progress by polling serial port output."""
+        _LOG.info("Waiting for customization script to finish and VM shutdown...")
+        time.sleep(15)  # Allow initial VM boot
+
+        offset = 0
+        start_time = time.time()
+        timeout_secs = 7200
+        if hasattr(args, "build_timeout_sec") and isinstance(
+            args.build_timeout_sec, (int, float)
+        ):
+            timeout_secs = args.build_timeout_sec
+        delay = 10
+
+        with open(local_log_file, "w") as log_f:
+            while time.time() - start_time < timeout_secs:
+                try:
+                    # Check VM instance state
+                    inst = self.instances_client.get(
+                        project=args.project_id,
+                        zone=args.zone,
+                        instance=instance_name,
+                        retry=_DEFAULT_RETRY,
+                    )
+                    _LOG.info("VM Status: %s", inst.status)
+                    is_stopped = inst.status in ("TERMINATED", "STOPPED")
+
+                    # Retrieve serial port output
+                    try:
+                        res = self.instances_client.get_serial_port_output(
+                            request={
+                                "project": args.project_id,
+                                "zone": args.zone,
+                                "instance": instance_name,
+                                "port": 1,
+                                "start": offset,
+                            },
+                            retry=_DEFAULT_RETRY,
+                        )
+                        if res.contents:
+                            log_f.write(res.contents)
+                            log_f.flush()
+                            sys.stdout.write(res.contents)
+                            sys.stdout.flush()
+                            offset = res.next_
+
+                            # Reset backoff delay when new logs arrive (Point 7)
+                            delay = 10
+
+                            if _has_build_signal(res.contents, "BuildSucceeded:"):
+                                state.build_succeeded = True
+                                _LOG.info("Customization script succeeded.")
+                            elif _has_build_signal(res.contents, "BuildFailed:"):
+                                state.build_failed = True
+                                _LOG.info("Customization script failed.")
+                        else:
+                            # Apply exponential backoff when no output is received
+                            delay = min(delay * 2, 60)
+
+                    except (
+                        GoogleAPIError,
+                        NotFound,
+                        PermissionDenied,
+                        Conflict,
+                        DeadlineExceeded,
+                    ) as e:
+                        if is_stopped:
+                            _LOG.info(
+                                "VM is stopped and serial output is no longer available: %s",
+                                e,
+                            )
+                        else:
+                            raise
+
+                    if is_stopped or state.build_succeeded or state.build_failed:
+                        break
+
+                    time.sleep(delay)
+                except (
+                    GoogleAPIError,
+                    NotFound,
+                    PermissionDenied,
+                    Conflict,
+                    DeadlineExceeded,
+                ) as e:
+                    _LOG.warning("Error reading serial output (will retry): %s", e)
+                    time.sleep(delay)
+
+        # Check outcome status in log file
+        if not state.build_succeeded:
+            with open(local_log_file, "r") as log_f:
+                logs = log_f.read()
+                if _has_build_signal(logs, "BuildSucceeded:"):
+                    state.build_succeeded = True
+                elif _has_build_signal(logs, "BuildFailed:"):
+                    state.build_failed = True
+
+    def _create_final_image(self, args, disk_name):
+        """Creates the final custom image from GCE boot disk."""
+        _LOG.info("Creating custom image %s from disk...", args.image_name)
+        image_resource = compute_v1.Image(
+            name=args.image_name,
+            source_disk=f"projects/{args.project_id}/zones/{args.zone}/disks/{disk_name}",
+            family=args.family,
+        )
+        if args.storage_location:
+            image_resource.storage_locations = [args.storage_location]
+
+        try:
+            op = self.images_client.insert(
+                project=args.project_id,
+                image_resource=image_resource,
+                retry=_DEFAULT_RETRY,
+            )
+            self.compute_helper.wait_for_global_operation(args.project_id, op.name)
+            _LOG.info("Successfully created custom image %s.", args.image_name)
+        except (
+            GoogleAPIError,
+            NotFound,
+            PermissionDenied,
+            Conflict,
+            DeadlineExceeded,
+        ) as e:
+            raise RuntimeError(
+                f"Error creating final custom image {args.image_name}: {e}"
+            )
+
+    def _cleanup(self, args, instance_name, disk_name, state):
+        """Deletes provisioned VM and disk on completion or failure."""
+        if state.vm_created:
+            try:
+                _LOG.info("Cleaning up VM instance %s...", instance_name)
+                op = self.instances_client.delete(
+                    project=args.project_id,
+                    zone=args.zone,
+                    instance=instance_name,
+                    retry=_DEFAULT_RETRY,
+                )
+                self.compute_helper.wait_for_zone_operation(
+                    args.project_id, args.zone, op.name
+                )
+            except Exception as e:
+                _LOG.warning("Failed to delete VM instance %s: %s", instance_name, e)
+
+        if state.disk_created and not state.vm_created:
+            try:
+                _LOG.info("Cleaning up boot disk %s...", disk_name)
+                op = self.disks_client.delete(
+                    project=args.project_id,
+                    zone=args.zone,
+                    disk=disk_name,
+                    retry=_DEFAULT_RETRY,
+                )
+                self.compute_helper.wait_for_zone_operation(
+                    args.project_id, args.zone, op.name
+                )
+            except NotFound:
+                _LOG.info("Boot disk %s was already deleted (auto-delete).", disk_name)
+            except Exception as e:
+                _LOG.warning("Failed to delete boot disk %s: %s", disk_name, e)
+
+    def _upload_logs(self, args, bucket):
+        """Syncs local logs to GCS."""
+        try:
+            _LOG.info("Syncing local logs to GCS log folder...")
+            for root, _, files in os.walk(args.log_dir):
+                for name in files:
+                    path = os.path.join(root, name)
+                    rel = os.path.relpath(path, args.log_dir)
+                    blob_path = f"{args.gcs_base_path}/logs/{rel}"
+                    blob = bucket.blob(blob_path)
+                    blob.upload_from_filename(path, retry=_DEFAULT_RETRY)
+        except (
+            GoogleAPIError,
+            NotFound,
+            PermissionDenied,
+            Conflict,
+            DeadlineExceeded,
+        ) as e:
+            _LOG.warning("Failed to upload build logs to GCS: %s", e)
+
+    def add_label(self, args):
+        """Sets Dataproc version label in the custom image."""
+        if args.dry_run:
+            _LOG.info("Dry-run mode: Skipping label attachment.")
+            return
+
+        _LOG.info("Setting label on custom image via API client...")
+        version_label = args.dataproc_version.replace(".", "-").lower()
+
+        # Retrieve current image to get resource fingerprint
+        img = self.images_client.get(
+            project=args.project_id, image=args.image_name, retry=_DEFAULT_RETRY
+        )
+
+        # Set labels
+        labels_spec = compute_v1.GlobalSetLabelsRequest(
+            label_fingerprint=img.label_fingerprint,
+            labels={"goog-dataproc-version": version_label},
+        )
+        op = self.images_client.set_labels(
+            project=args.project_id,
+            resource=args.image_name,
+            global_set_labels_request_resource=labels_spec,
+            retry=_DEFAULT_RETRY,
+        )
+        self.compute_helper.wait_for_global_operation(args.project_id, op.name)
+        _LOG.info("Successfully set label on custom image %s.", args.image_name)
+
+    def notify_expiration(self, args):
+        """Notifies when the image will expire using Images API client."""
+        if args.dry_run:
+            _LOG.info("Dry-run mode: Skipping expiration notification.")
+            return
+
+        _LOG.info("Successfully built Dataproc custom image: %s", args.image_name)
+        img = self.images_client.get(
+            project=args.project_id, image=args.image_name, retry=_DEFAULT_RETRY
+        )
+        timestamp_string = img.creation_timestamp
+
+        # RFC3339 timestamp parsing
+        creation_date = datetime.datetime.fromisoformat(
+            timestamp_string.replace("Z", "+00:00")
+        )
+        expiration_date = creation_date + datetime.timedelta(days=365)
+
+        notification_text = """
+#####################################################################
+  WARNING: DATAPROC CUSTOM IMAGE '{}'
+           WILL EXPIRE ON {}.
+#####################################################################
+"""
+        _LOG.warning(notification_text.format(args.image_name, str(expiration_date)))
+
+    def run_smoke_test(self, args):
+        """Runs a smoke test on the custom image."""
+        from custom_image_utils import smoke_test_runner
+
+        smoke_test_runner.run(args)
diff --git a/custom_image_utils/args_inferer.py b/custom_image_utils/args_inferer.py
index 75bfc09..4b1507b 100644
--- a/custom_image_utils/args_inferer.py
+++ b/custom_image_utils/args_inferer.py
@@ -228,6 +228,8 @@ def _infer_base_image(args):
         "Neither --dataproc-version nor --base-image-uri nor --source-image-family-uri is specified.")
   _LOG.info("Returned Dataproc base image: %s", args.dataproc_base_image)
   _LOG.info("Returned Dataproc version   : %s", args.dataproc_version)
+  print("Returned Dataproc base image: {}".format(args.dataproc_base_image))
+  print("Returned Dataproc version   : {}".format(args.dataproc_version))
 
 
 def _infer_oauth(args):
diff --git a/custom_image_utils/args_parser.py b/custom_image_utils/args_parser.py
index ebd1aee..d2068ba 100644
--- a/custom_image_utils/args_parser.py
+++ b/custom_image_utils/args_parser.py
@@ -29,7 +29,7 @@
 _VERSION_REGEX = re.compile(r"^\d+\.\d+\.\d+(-RC\d+)?(-[a-z\-]+\d+)?$")
 _FULL_IMAGE_URI = re.compile(r"^(https://www\.googleapis\.com/compute/([^/]+)/)?projects/([^/]+)/global/images/([^/]+)$")
 _FULL_IMAGE_FAMILY_URI = re.compile(r"^(https://www\.googleapis\.com/compute/([^/]+)/)?projects/([^/]+)/global/images/family/([^/]+)$")
-_LATEST_FROM_MINOR_VERSION = re.compile(r"^(\d+)\.(\d+)-((?:debian|ubuntu|rocky)\d+)$")
+_LATEST_FROM_MINOR_VERSION = re.compile(r"^(\d+)\.(\d+)-(?:[a-zA-Z0-9\-]+-)?((?:debian|ubuntu|rocky|centos)\d+)$")
 _ARM_ARCH_REGEX = re.compile(r"""
   (?:^|[-_./])  # Non-alphanumeric separator or start of string
   (?:           # Match one of the ARM architecture terms:
@@ -264,6 +264,16 @@ def parse_args(args):
       default="googleapis.com",
       help="""(Optional) The universe domain to configure for gcloud. Defaults to 'googleapis.com'."""
   )
+  parser.add_argument(
+      "--execution-engine",
+      type=str,
+      required=False,
+      choices=["cli", "api"],
+      default="cli",
+      help="""(Optional) The execution engine used to create the image.
+      Defaults to 'cli' (gcloud / gsutil command line tools).
+      Use 'api' to use native Google Cloud client libraries.
+      """)
 
   parsed_args = parser.parse_args(args)
 
diff --git a/custom_image_utils/cli_execution_engine.py b/custom_image_utils/cli_execution_engine.py
new file mode 100644
index 0000000..3c6b2d3
--- /dev/null
+++ b/custom_image_utils/cli_execution_engine.py
@@ -0,0 +1,57 @@
+# Copyright 2026 Google LLC and contributors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#            http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""CLI-based implementation of the ExecutionEngine interface."""
+
+import subprocess
+from custom_image_utils import args_inferer
+from custom_image_utils import expiration_notifier
+from custom_image_utils import image_labeller
+from custom_image_utils import shell_image_creator
+from custom_image_utils import smoke_test_runner
+from custom_image_utils.execution_engine import ExecutionEngine
+
+
+class CliExecutionEngine(ExecutionEngine):
+    """Execution engine that uses gcloud and gsutil CLI utilities."""
+
+    def infer_args(self, args):
+        args_inferer.infer_args(args)
+
+    def perform_sanity_checks(self, args):
+        # Check the image doesn't already exist using gcloud compute images describe.
+        command = [
+            "gcloud",
+            "compute",
+            "images",
+            "describe",
+            args.image_name,
+            f"--project={args.project_id}",
+        ]
+        result = subprocess.run(
+            command, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
+        )
+        if result.returncode == 0:
+            raise RuntimeError("Image {} already exists.".format(args.image_name))
+
+    def create_image(self, args):
+        shell_image_creator.create(args)
+
+    def add_label(self, args):
+        image_labeller.add_label(args)
+
+    def run_smoke_test(self, args):
+        smoke_test_runner.run(args)
+
+    def notify_expiration(self, args):
+        expiration_notifier.notify(args)
diff --git a/custom_image_utils/compute_operation_helper.py b/custom_image_utils/compute_operation_helper.py
new file mode 100644
index 0000000..6518163
--- /dev/null
+++ b/custom_image_utils/compute_operation_helper.py
@@ -0,0 +1,88 @@
+# Copyright 2026 Google LLC and contributors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#            http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Helper for waiting for Compute Engine long-running operations."""
+
+import logging
+import time
+from google.cloud import compute_v1
+
+_LOG = logging.getLogger(__name__)
+
+
+class ComputeOperationHelper:
+    """Helper to wait for Google Compute Engine long-running operations."""
+
+    def __init__(self, credentials=None):
+        self.zone_client = compute_v1.ZoneOperationsClient(credentials=credentials)
+        self.global_client = compute_v1.GlobalOperationsClient(credentials=credentials)
+        self.region_client = compute_v1.RegionOperationsClient(credentials=credentials)
+
+    def wait_for_zone_operation(self, project, zone, operation_name, timeout_secs=600):
+        """Waits for a zonal operation to complete."""
+        _LOG.info("Waiting for zonal operation %s in zone %s...", operation_name, zone)
+        start_time = time.time()
+        while time.time() - start_time < timeout_secs:
+            op = self.zone_client.get(
+                project=project, zone=zone, operation=operation_name
+            )
+            if op.status == compute_v1.Operation.Status.DONE:
+                if op.error:
+                    raise RuntimeError("Zonal operation failed: {}".format(op.error))
+                return op
+            time.sleep(5)
+        raise TimeoutError(
+            "Zonal operation {} timed out after {}s".format(
+                operation_name, timeout_secs
+            )
+        )
+
+    def wait_for_global_operation(self, project, operation_name, timeout_secs=600):
+        """Waits for a global operation to complete."""
+        _LOG.info("Waiting for global operation %s...", operation_name)
+        start_time = time.time()
+        while time.time() - start_time < timeout_secs:
+            op = self.global_client.get(project=project, operation=operation_name)
+            if op.status == compute_v1.Operation.Status.DONE:
+                if op.error:
+                    raise RuntimeError("Global operation failed: {}".format(op.error))
+                return op
+            time.sleep(5)
+        raise TimeoutError(
+            "Global operation {} timed out after {}s".format(
+                operation_name, timeout_secs
+            )
+        )
+
+    def wait_for_region_operation(
+        self, project, region, operation_name, timeout_secs=600
+    ):
+        """Waits for a regional operation to complete."""
+        _LOG.info(
+            "Waiting for regional operation %s in region %s...", operation_name, region
+        )
+        start_time = time.time()
+        while time.time() - start_time < timeout_secs:
+            op = self.region_client.get(
+                project=project, region=region, operation=operation_name
+            )
+            if op.status == compute_v1.Operation.Status.DONE:
+                if op.error:
+                    raise RuntimeError("Regional operation failed: {}".format(op.error))
+                return op
+            time.sleep(5)
+        raise TimeoutError(
+            "Regional operation {} timed out after {}s".format(
+                operation_name, timeout_secs
+            )
+        )
diff --git a/custom_image_utils/execution_engine.py b/custom_image_utils/execution_engine.py
new file mode 100644
index 0000000..e3f7c21
--- /dev/null
+++ b/custom_image_utils/execution_engine.py
@@ -0,0 +1,50 @@
+# Copyright 2026 Google LLC and contributors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#            http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Base execution engine interface for custom image creation workflow."""
+
+import abc
+
+
+class ExecutionEngine(abc.ABC):
+    """Abstract base class representing the execution engine interface."""
+
+    @abc.abstractmethod
+    def infer_args(self, args):
+        """Infers missing command-line arguments using Google Cloud APIs or CLI."""
+        pass
+
+    @abc.abstractmethod
+    def perform_sanity_checks(self, args):
+        """Performs sanity checks (e.g. checks if the target image already exists)."""
+        pass
+
+    @abc.abstractmethod
+    def create_image(self, args):
+        """Creates the custom Dataproc image."""
+        pass
+
+    @abc.abstractmethod
+    def add_label(self, args):
+        """Adds the goog-dataproc-version label to the created custom image."""
+        pass
+
+    @abc.abstractmethod
+    def run_smoke_test(self, args):
+        """Runs a smoke test on the custom image."""
+        pass
+
+    @abc.abstractmethod
+    def notify_expiration(self, args):
+        """Notifies the user when the custom image will expire."""
+        pass
diff --git a/custom_image_utils/shell_script_generator.py b/custom_image_utils/shell_script_generator.py
index 733bce9..fdd5fb0 100644
--- a/custom_image_utils/shell_script_generator.py
+++ b/custom_image_utils/shell_script_generator.py
@@ -287,6 +287,10 @@ class Generator:
 
   def _init_args(self, args):
     self.args = args
+    self.args.setdefault("optional_components", None)
+    self.args.setdefault("dataproc_version", None)
+    self.args.setdefault("universe_domain", "googleapis.com")
+    self.args.setdefault("trusted_cert", "tls/db.der")
     if "run_id" not in self.args:
       self.args["run_id"] = "custom-image-{image_name}-{timestamp}".format(
           timestamp=datetime.now().strftime("%Y%m%d-%H%M%S"), **self.args)
diff --git a/generate_custom_image.py b/generate_custom_image.py
index 6f2e32d..614ec0e 100644
--- a/generate_custom_image.py
+++ b/generate_custom_image.py
@@ -34,61 +34,61 @@
 
 import logging
 import os
-import subprocess
 import sys
 
-from custom_image_utils import args_inferer
 from custom_image_utils import args_parser
-from custom_image_utils import expiration_notifier
-from custom_image_utils import image_labeller
-from custom_image_utils import shell_image_creator
-from custom_image_utils import smoke_test_runner
 
 logging.basicConfig()
 _LOG = logging.getLogger(__name__)
 _LOG.setLevel(logging.WARN)
 
 
-def parse_args(raw_args):
-  """Parses and infers command line arguments."""
-
-  args = args_parser.parse_args(raw_args)
-  _LOG.info("Parsed args: {}".format(args))
-  args_inferer.infer_args(args)
-  _LOG.info("Inferred args: {}".format(args))
-  return args
+def get_execution_engine(args):
+  """Instantiates the appropriate execution engine."""
+  if args.execution_engine == "api":
+    from custom_image_utils.api_execution_engine import ApiExecutionEngine
 
+    credentials = None
+    if args.oauth:
+      import google.auth
 
-def perform_sanity_checks(args):
-  _LOG.info("Performing sanity checks...")
+      credentials, _ = google.auth.load_credentials_from_file(args.oauth)
 
-  # Customization script
-  if not os.path.isfile(args.customization_script):
-    raise Exception("Invalid path to customization script: '{}' is not a file.".format(
-        args.customization_script))
+    return ApiExecutionEngine(credentials=credentials)
+  else:
+    from custom_image_utils.cli_execution_engine import CliExecutionEngine
 
-  # Check the image doesn't already exist.
-  command = "gcloud compute images describe {} --project={}".format(
-      args.image_name, args.project_id)
-  with open(os.devnull, 'w') as devnull:
-    pipe = subprocess.Popen(
-        [command], stdout=devnull, stderr=devnull, shell=True)
-    pipe.wait()
-    if pipe.returncode == 0:
-      raise RuntimeError("Image {} already exists.".format(args.image_name))
-
-  _LOG.info("Passed sanity checks...")
+    return CliExecutionEngine()
 
 
 def main():
   """Generates custom image."""
 
-  args = parse_args(sys.argv[1:])
-  perform_sanity_checks(args)
-  shell_image_creator.create(args)
-  image_labeller.add_label(args)
-  smoke_test_runner.run(args)
-  expiration_notifier.notify(args)
+  # Parse args
+  args = args_parser.parse_args(sys.argv[1:])
+  _LOG.info("Parsed args: {}".format(args))
+
+  # Get selected execution engine
+  engine = get_execution_engine(args)
+
+  # Infer remaining arguments and check customization script path
+  is_gcs_script = args.customization_script.startswith("gs://")
+  if not is_gcs_script and not os.path.isfile(args.customization_script):
+    raise Exception(
+        "Invalid path to customization script: '{}' is not a file.".format(
+            args.customization_script
+        )
+    )
+
+  engine.infer_args(args)
+  _LOG.info("Inferred args: {}".format(args))
+
+  # Run custom image creation workflow
+  engine.perform_sanity_checks(args)
+  engine.create_image(args)
+  engine.add_label(args)
+  engine.run_smoke_test(args)
+  engine.notify_expiration(args)
 
 
 if __name__ == "__main__":
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..dc54a3e
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,5 @@
+# Requirements for API-based custom image generation
+google-cloud-compute
+google-cloud-storage
+google-auth
+ruff
diff --git a/scripts/test_customization.sh b/scripts/test_customization.sh
new file mode 100644
index 0000000..bd9b1aa
--- /dev/null
+++ b/scripts/test_customization.sh
@@ -0,0 +1,23 @@
+#!/bin/bash
+# Copyright 2026 Google LLC and contributors
+# Licensed under the Apache License, Version 2.0.
+#
+# Simple customization script for testing the Dataproc custom image build process.
+
+set -eo pipefail
+
+echo "========================================="
+echo "Starting local test customization script..."
+echo "========================================="
+
+# 1. Update package list and install a lightweight diagnostic tool (e.g., htop or tree)
+echo "Installing htop utility..."
+apt-get update && apt-get install -y htop
+
+# 2. Write a verification marker file to the image
+echo "Writing verification marker file to /etc/custom_image_test..."
+echo "Dataproc Custom Image built successfully at $(date)" > /etc/custom_image_test
+
+echo "========================================="
+echo "Customization script finished successfully!"
+echo "========================================="
diff --git a/tests/test_api_execution_engine.py b/tests/test_api_execution_engine.py
new file mode 100644
index 0000000..c766e2a
--- /dev/null
+++ b/tests/test_api_execution_engine.py
@@ -0,0 +1,309 @@
+# Copyright 2026 Google LLC and contributors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#            http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Unit tests for api_execution_engine.py."""
+
+import unittest
+from unittest.mock import MagicMock, patch, mock_open
+
+from custom_image_utils.api_execution_engine import ApiExecutionEngine
+
+
+class TestApiExecutionEngine(unittest.TestCase):
+    def setUp(self):
+        self.auth_patch = patch(
+            "google.auth.default", return_value=(None, "test-project")
+        )
+        self.auth_patch.start()
+        self.addCleanup(self.auth_patch.stop)
+
+        self.images_patch = patch("google.cloud.compute_v1.ImagesClient")
+        self.images_patch.start()
+        self.addCleanup(self.images_patch.stop)
+
+        self.disks_patch = patch("google.cloud.compute_v1.DisksClient")
+        self.disks_patch.start()
+        self.addCleanup(self.disks_patch.stop)
+
+        self.instances_patch = patch("google.cloud.compute_v1.InstancesClient")
+        self.instances_patch.start()
+        self.addCleanup(self.instances_patch.stop)
+
+        self.storage_patch = patch("google.cloud.storage.Client")
+        self.storage_patch.start()
+        self.addCleanup(self.storage_patch.stop)
+
+        self.engine = ApiExecutionEngine()
+        self.engine.compute_helper = MagicMock()
+
+    def test_get_default_project(self):
+        with patch("google.auth.default", return_value=(None, "test-project")):
+            project_id = self.engine._get_default_project()
+            self.assertEqual(project_id, "test-project")
+
+    def test_get_default_project_with_credentials(self):
+        mock_creds = MagicMock()
+        mock_creds.project_id = "credentials-project"
+        self.engine.credentials = mock_creds
+        try:
+            project_id = self.engine._get_default_project()
+            self.assertEqual(project_id, "credentials-project")
+        finally:
+            self.engine.credentials = None
+
+    def test_get_default_project_missing(self):
+        with patch("google.auth.default", return_value=(None, None)):
+            with self.assertRaises(RuntimeError):
+                self.engine._get_default_project()
+
+    def test_infer_args_with_project_and_base_image(self):
+        args = MagicMock()
+        args.project_id = "test-project"
+        args.base_image_uri = (
+            "projects/cloud-dataproc/global/images/dataproc-2-1-deb11-20260611"
+        )
+        args.base_image_family = None
+        args.dataproc_version = None
+        args.oauth = None
+        args.network = None
+        args.subnetwork = None
+        args.zone = "us-central1-a"
+        args.shutdown_instance_timer_sec = 300
+
+        mock_image = MagicMock()
+        mock_image.labels = {"goog-dataproc-version": "2.1.115-debian11"}
+        self.engine.images_client.get.return_value = mock_image
+
+        self.engine.infer_args(args)
+
+        self.assertEqual(
+            args.dataproc_base_image,
+            "projects/cloud-dataproc/global/images/dataproc-2-1-deb11-20260611",
+        )
+        self.assertEqual(args.dataproc_version, "2.1.115-debian11")
+        self.assertEqual(args.network, "projects/test-project/global/networks/default")
+
+    def test_perform_sanity_checks_not_found(self):
+        from google.api_core.exceptions import NotFound
+
+        args = MagicMock()
+        args.project_id = "test-project"
+        args.image_name = "test-image"
+
+        self.engine.images_client.get.side_effect = NotFound("Not Found")
+        # Should not raise exception
+        self.engine.perform_sanity_checks(args)
+
+    def test_perform_sanity_checks_already_exists(self):
+        args = MagicMock()
+        args.project_id = "test-project"
+        args.image_name = "test-image"
+
+        self.engine.images_client.get.return_value = MagicMock()
+        with self.assertRaises(RuntimeError) as ctx:
+            self.engine.perform_sanity_checks(args)
+        self.assertIn("already exists", str(ctx.exception))
+
+    @patch("builtins.open", new_callable=mock_open, read_data="echo test")
+    @patch("os.makedirs")
+    @patch("time.sleep")
+    def test_create_image_dry_run(self, mock_sleep, mock_makedirs, mock_open):
+        args = MagicMock()
+        args.dry_run = True
+
+        # Should exit early without making any API calls
+        self.engine.create_image(args)
+        self.engine.images_client.insert.assert_not_called()
+
+    def test_create_disk_success(self):
+        args = MagicMock()
+        args.project_id = "test-project"
+        args.zone = "us-central1-a"
+        args.dataproc_base_image = "projects/cloud-dataproc/global/images/dataproc-2-1"
+        args.disk_size = 50
+
+        self.engine.disks_client.insert.return_value = MagicMock(name="op")
+        self.engine._create_disk(args, "test-disk")
+        self.engine.disks_client.insert.assert_called_once()
+
+    def test_create_disk_failure(self):
+        from google.api_core.exceptions import GoogleAPIError
+
+        args = MagicMock()
+        args.project_id = "test-project"
+        args.zone = "us-central1-a"
+        args.dataproc_base_image = "projects/cloud-dataproc/global/images/dataproc-2-1"
+        args.disk_size = 50
+
+        self.engine.disks_client.insert.side_effect = GoogleAPIError("GCP error")
+        with self.assertRaises(RuntimeError) as ctx:
+            self.engine._create_disk(args, "test-disk")
+        self.assertIn("Error creating boot disk", str(ctx.exception))
+
+    @patch("builtins.open", new_callable=mock_open, read_data="echo startup")
+    def test_create_vm_success(self, mock_open):
+        args = MagicMock()
+        args.project_id = "test-project"
+        args.zone = "us-central1-a"
+        args.machine_type = "n1-standard-4"
+        args.service_account = "test-sa@project.iam.gserviceaccount.com"
+        args.subnetwork = None
+        args.network = "global/networks/default"
+        args.no_external_ip = False
+        args.accelerator = None
+        args.metadata = "key1=value1"
+        args.universe_domain = "googleapis.com"
+        args.dataproc_version = "2.1.115"
+        args.custom_sources_path = "gs://test-bucket/sources"
+        args.shutdown_timer_in_sec = 300
+
+        self.engine.instances_client.insert.return_value = MagicMock(name="op")
+        self.engine._create_vm(args, "test-instance", "test-disk", "us-central1")
+        self.engine.instances_client.insert.assert_called_once()
+
+    @patch("builtins.open", new_callable=mock_open, read_data="BuildSucceeded:")
+    @patch("time.sleep")
+    def test_monitor_build_success(self, mock_sleep, mock_open_file):
+        from custom_image_utils.api_execution_engine import BuildState
+
+        args = MagicMock()
+        args.project_id = "test-project"
+        args.zone = "us-central1-a"
+        state = BuildState()
+
+        mock_instance = MagicMock()
+        mock_instance.status = "RUNNING"
+        self.engine.instances_client.get.return_value = mock_instance
+
+        mock_serial = MagicMock()
+        mock_serial.contents = "startup-script: BuildSucceeded:\n"
+        mock_serial.next_ = 100
+        self.engine.instances_client.get_serial_port_output.return_value = mock_serial
+
+        self.engine._monitor_build(args, "test-instance", "local-log.log", state)
+        self.assertTrue(state.build_succeeded)
+
+    def test_create_final_image_success(self):
+        args = MagicMock()
+        args.project_id = "test-project"
+        args.image_name = "test-image"
+        args.family = "test-family"
+        args.storage_location = "us"
+
+        self.engine.images_client.insert.return_value = MagicMock(name="op")
+        self.engine._create_final_image(args, "test-disk")
+        self.engine.images_client.insert.assert_called_once()
+
+    def test_cleanup_vm_and_disk(self):
+        from custom_image_utils.api_execution_engine import BuildState
+
+        args = MagicMock()
+        args.project_id = "test-project"
+        args.zone = "us-central1-a"
+
+        # Scenario 1: Both created -> Only VM delete is called (disk is auto-deleted with VM)
+        state = BuildState(disk_created=True, vm_created=True)
+        self.engine._cleanup(args, "test-instance", "test-disk", state)
+        self.engine.instances_client.delete.assert_called_once()
+        self.engine.disks_client.delete.assert_not_called()
+
+        # Scenario 2: Only disk created -> Only disk delete is called
+        self.engine.instances_client.delete.reset_mock()
+        self.engine.disks_client.delete.reset_mock()
+        state2 = BuildState(disk_created=True, vm_created=False)
+        self.engine._cleanup(args, "test-instance", "test-disk", state2)
+        self.engine.instances_client.delete.assert_not_called()
+        self.engine.disks_client.delete.assert_called_once()
+
+    def test_csv_metadata_parsing(self):
+        args = MagicMock()
+        args.project_id = "test-project"
+        args.zone = "us-central1-a"
+        args.machine_type = "n1-standard-4"
+        args.service_account = "test-sa@project.iam.gserviceaccount.com"
+        args.subnetwork = None
+        args.network = "global/networks/default"
+        args.no_external_ip = False
+        args.accelerator = None
+        args.metadata = 'key1="val1,val2",key2=val3'
+        args.universe_domain = "googleapis.com"
+        args.dataproc_version = "2.1.115"
+        args.custom_sources_path = "gs://test-bucket/sources"
+        args.shutdown_timer_in_sec = 300
+        args.bucket_name = "test-bucket"
+        args.gcs_base_path = "run-id"
+
+        self.engine.instances_client.insert.return_value = MagicMock(name="op")
+
+        with patch("builtins.open", mock_open(read_data="echo test")):
+            self.engine._create_vm(args, "test-instance", "test-disk", "us-central1")
+
+        call_args = self.engine.instances_client.insert.call_args
+        instance_resource = call_args.kwargs["instance_resource"]
+        metadata_items = instance_resource.metadata.items
+
+        metadata_dict = {item.key: item.value for item in metadata_items}
+        self.assertEqual(metadata_dict.get("key1"), "val1,val2")
+        self.assertEqual(metadata_dict.get("key2"), "val3")
+
+    def test_dataproc_version_parsing_and_sorting(self):
+        from custom_image_utils.api_execution_engine import _parse_dataproc_version
+
+        v1 = _parse_dataproc_version("2.1.115-debian11")
+        v2 = _parse_dataproc_version("2.0.50-debian10")
+        v3 = _parse_dataproc_version("2.1.99-debian11")
+        v4 = _parse_dataproc_version("2-1-115-debian11")
+
+        self.assertEqual(v1, (2, 1, 115))
+        self.assertEqual(v2, (2, 0, 50))
+        self.assertEqual(v3, (2, 1, 99))
+        self.assertEqual(v4, (2, 1, 115))
+        self.assertTrue(v1 > v3)
+        self.assertTrue(v3 > v2)
+
+    @patch("custom_image_utils.smoke_test_runner.run")
+    def test_run_smoke_test(self, mock_smoke_run):
+        args = MagicMock()
+        self.engine.run_smoke_test(args)
+        mock_smoke_run.assert_called_once_with(args)
+
+    def test_invalid_accelerator_count(self):
+        args = MagicMock()
+        args.project_id = "test-project"
+        args.zone = "us-central1-a"
+        args.machine_type = "n1-standard-4"
+        args.service_account = "test-sa@project.iam.gserviceaccount.com"
+        args.subnetwork = None
+        args.network = "global/networks/default"
+        args.no_external_ip = False
+        args.accelerator = "type=nvidia-tesla-v100,count=two"
+        args.universe_domain = "googleapis.com"
+        args.dataproc_version = "2.1.115"
+        args.custom_sources_path = "gs://test-bucket/sources"
+        args.shutdown_timer_in_sec = 300
+        args.bucket_name = "test-bucket"
+        args.gcs_base_path = "run-id"
+        args.metadata = None
+
+        with patch("builtins.open", mock_open(read_data="echo test")):
+            with self.assertRaises(RuntimeError) as ctx:
+                self.engine._create_vm(
+                    args, "test-instance", "test-disk", "us-central1"
+                )
+        self.assertIn(
+            "Invalid accelerator count: two. Must be an integer.", str(ctx.exception)
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/test_args_parser.py b/tests/test_args_parser.py
index e0cc4d4..c6b8934 100644
--- a/tests/test_args_parser.py
+++ b/tests/test_args_parser.py
@@ -208,6 +208,10 @@ def _args_exception(dataproc_version):
       raise e
 
   def _make_expected_result(self, **kwargs):
+    if "execution_engine" not in kwargs:
+      kwargs["execution_engine"] = "cli"
+    if "universe_domain" not in kwargs:
+      kwargs["universe_domain"] = "googleapis.com"
     return argparse.Namespace(**kwargs)
 
 if __name__ == '__main__':
diff --git a/tests/test_shell_script_generator.py b/tests/test_shell_script_generator.py
index d14253a..e414150 100644
--- a/tests/test_shell_script_generator.py
+++ b/tests/test_shell_script_generator.py
@@ -16,119 +16,259 @@
 
 from custom_image_utils import shell_script_generator
 
-_expected_script = """
+_expected_script = r"""
 #!/usr/bin/env bash
 
 # Script for creating Dataproc custom image.
 
-set -euxo pipefail
+set -euo pipefail
 
-RED='\\e[0;31m'
-GREEN='\\e[0;32m'
-NC='\\e[0m'
+RED='\e[0;31m'
+GREEN='\e[0;32m'
+NC='\e[0m'
+
+base_obj_type="images"
+
+function execute_with_retries() (
+  set +x
+  local -r cmd="$*"
+
+  for ((i = 0; i < 3; i++)); do
+    if eval "$cmd"; then return 0 ; fi
+    sleep 12
+  done
+  return 1
+)
+
+function prepare() {
+  # With the 402.0.0 release of gcloud sdk, `gcloud storage` can be
+  # used as a more performant replacement for `gsutil`
+  if gcloud --help >/dev/null 2>&1 && gcloud storage --help >/dev/null 2>&1; then
+    gsutil_cmd="gcloud storage"
+    rsync_cmd="${gsutil_cmd} rsync"
+  else
+    gsutil_cmd="gsutil -o GSUtil:check_hashes=never"
+    rsync_cmd="${gsutil_cmd} -m rsync"
+  fi
+}
 
 function exit_handler() {
   echo 'Cleaning up before exiting.'
 
   if [[ -f /tmp/custom-image-my-image-20190611-160823/vm_created ]]; then
     echo 'Deleting VM instance.'
-    gcloud compute instances delete my-image-install         --project=my-project --zone=us-west1-a -q
+    execute_with_retries       gcloud compute instances delete my-image-install --project=my-project --zone=us-west1-a -q
   elif [[ -f /tmp/custom-image-my-image-20190611-160823/disk_created ]]; then
     echo 'Deleting disk.'
-    gcloud compute disks delete my-image-install --project=my-project --zone=us-west1-a -q
+    execute_with_retries gcloud compute ${base_obj_type} delete my-image-install --project=my-project -q
   fi
 
   echo 'Uploading local logs to GCS bucket.'
-  gcloud storage rsync --recursive /tmp/custom-image-my-image-20190611-160823/logs/ gs://my-bucket/custom-image-my-image-20190611-160823/logs/
+  ${rsync_cmd} -r /tmp/custom-image-my-image-20190611-160823/logs/ gs://my-bucket/custom-image-my-image-20190611-160823/logs/
 
   if [[ -f /tmp/custom-image-my-image-20190611-160823/image_created ]]; then
-    echo -e "${GREEN}Workflow succeeded, check logs at /tmp/custom-image-my-image-20190611-160823/logs/ or gs://my-bucket/custom-image-my-image-20190611-160823/logs/${NC}"
+    echo -e "${GREEN}Workflow succeeded${NC}, check logs at /tmp/custom-image-my-image-20190611-160823/logs/ or gs://my-bucket/custom-image-my-image-20190611-160823/logs/"
     exit 0
   else
-    echo -e "${RED}Workflow failed, check logs at /tmp/custom-image-my-image-20190611-160823/logs/ or gs://my-bucket/custom-image-my-image-20190611-160823/logs/${NC}"
+    echo -e "${RED}Workflow failed${NC}, check logs at /tmp/custom-image-my-image-20190611-160823/logs/ or gs://my-bucket/custom-image-my-image-20190611-160823/logs/"
     exit 1
   fi
 }
 
+function test_element_in_array {
+  local test_element="$1" ; shift
+  local -a test_array=("$@")
+
+  for item in "${test_array[@]}"; do
+    if [[ "${item}" == "${test_element}" ]]; then return 0 ; fi
+  done
+  return 1
+}
+
+function print_modulus_md5sum {
+  local derfile="$1"
+  openssl x509 -noout -modulus -in "${derfile}" | openssl md5 | awk '{print $2}'
+}
+
+function print_img_dbs_modulus_md5sums() {
+  local long_img_name="$1"
+  local img_name="$(echo ${long_img_name} | sed -e 's:^.*/::')"
+  local json_tmpfile="/tmp/custom-image-my-image-20190611-160823/${img_name}.json"
+  gcloud compute images describe ${long_img_name} --format json > "${json_tmpfile}"
+
+  local -a db_certs=()
+  mapfile -t db_certs < <( cat ${json_tmpfile} | jq -r 'try .shieldedInstanceInitialState.dbs[].content' )
+
+  local -a modulus_md5sums=()
+  for key in "${!db_certs[@]}" ; do
+    local derfile="/tmp/custom-image-my-image-20190611-160823/${img_name}.${key}.der"
+    echo "${db_certs[${key}]}" |       perl -M'MIME::Base64(decode_base64url)' -ne 'chomp; print( decode_base64url($_) )'       > "${derfile}"
+    modulus_md5sums+=( $(print_modulus_md5sum "${derfile}") )
+  done
+
+  echo "${modulus_md5sums[@]}"
+}
+
 function main() {
   echo 'Uploading files to GCS bucket.'
-  declare -a sources_k=([0]='run.sh' [1]='init_actions.sh' [2]='ext'\\''ra_src.txt')
-  declare -a sources_v=([0]='startup_script/run.sh' [1]='/tmp/my-script.sh' [2]='/path/to/extra.txt')
+  declare -a sources_k=([0]='run.sh' [1]='init_actions.sh' [2]='gce-proxy-setup.sh' [3]='ext'\''ra_src.txt')
+  declare -a sources_v=([0]='startup_script/run.sh' [1]='/tmp/my-script.sh' [2]='startup_script/gce-proxy-setup.sh' [3]='/path/to/extra.txt')
   for i in "${!sources_k[@]}"; do
-    gcloud storage cp "${sources_v[i]}" "gs://my-bucket/custom-image-my-image-20190611-160823/sources/${sources_k[i]}"
+    ${gsutil_cmd} cp "${sources_v[i]}" "gs://my-bucket/custom-image-my-image-20190611-160823/sources/${sources_k[i]}" > /dev/null 2>&1
   done
 
-  echo 'Creating disk.'
-  if [[ 'projects/my-dataproc-project/global/images/family/debian-10' = '' ||  'projects/my-dataproc-project/global/images/family/debian-10' = 'None' ]]; then
-     IMAGE_SOURCE="--image=projects/cloud-dataproc/global/images/dataproc-1-4-deb9-20190510-000000-rc01"
-  else
-     IMAGE_SOURCE="--image-family=projects/my-dataproc-project/global/images/family/debian-10"
+  local cert_args=""
+  local num_src_certs="0"
+  if [[ -n '' ]] && [[ -f '' ]]; then
+    # build tls/ directory from variables defined near the header of
+    # the examples/secure-boot/create-key-pair.sh file
+
+    eval "$(bash examples/secure-boot/create-key-pair.sh)"
+
+    # by default, a gcloud secret with the name of efi-db-pub-key-042 is
+    # created in the current project to store the certificate installed
+    # as the signature database file for this disk image
+
+    # The MS UEFI CA is a reasonable base from which to build trust.  We
+    # will trust code signed by this CA as well as code signed by
+    # trusted_cert (tls/db.der)
+
+    # The Microsoft Corporation UEFI CA 2011
+    local -r MS_UEFI_CA="tls/MicCorUEFCA2011_2011-06-27.crt"
+    test -f "${MS_UEFI_CA}" ||       curl -L -o ${MS_UEFI_CA} 'https://go.microsoft.com/fwlink/p/?linkid=321194'
+
+    local -a cert_list=()
+
+    local -a default_cert_list
+    default_cert_list=("" "${MS_UEFI_CA}")
+    local -a src_img_modulus_md5sums=()
+
+    mapfile -t src_img_modulus_md5sums < <(print_img_dbs_modulus_md5sums projects/cloud-dataproc/global/images/dataproc-1-4-deb9-20190510-000000-rc01)
+    num_src_certs="${#src_img_modulus_md5sums[@]}"
+    echo "${num_src_certs} db certificates attached to source image"
+    if [[ "${num_src_certs}" -eq "0" ]]; then
+      echo "no db certificates in source image"
+      cert_list=( "${default_cert_list[@]}" )
+    else
+      echo "db certs exist in source image"
+      for cert in ${default_cert_list[*]}; do
+        if test_element_in_array "$(print_modulus_md5sum ${cert})" ${src_img_modulus_md5sums[@]} ; then
+          echo "cert ${cert} is already in source image's db list"
+        else
+          cert_list+=("${cert}")
+        fi
+      done
+      # append source image's cert list
+      local img_name="$(echo projects/cloud-dataproc/global/images/dataproc-1-4-deb9-20190510-000000-rc01 | sed -e 's:^.*/::')"
+      if [[ ${#cert_list[@]} -ne 0 ]] && compgen -G "/tmp/custom-image-my-image-20190611-160823/${img_name}.*.der" > /dev/null ; then
+        cert_list+=(/tmp/custom-image-my-image-20190611-160823/${img_name}.*.der)
+      fi
+    fi
+
+    if [[ ${#cert_list[@]} -eq 0 ]]; then
+      echo "all certificates already included in source image's db list"
+    else
+      cert_args="--signature-database-file=$(IFS=, ; echo "${cert_list[*]}") --guest-os-features=UEFI_COMPATIBLE"
+    fi
   fi
-  
-  gcloud compute disks create my-image-install       --project=my-project       --zone=us-west1-a       ${IMAGE_SOURCE}       --type=pd-ssd       --size=40GB
 
-  touch "/tmp/custom-image-my-image-20190611-160823/disk_created"
+  date
+
+  if [[ -z "${cert_args}" && "${num_src_certs}" -ne "0" ]]; then
+    echo 'Re-using base image'
+    base_obj_type="reuse"
+    instance_disk_args='--image-project=my-project --image=projects/cloud-dataproc/global/images/dataproc-1-4-deb9-20190510-000000-rc01 --boot-disk-size=40G --boot-disk-type=pd-ssd'
+
+  elif [[ -n "${cert_args}" ]] ; then
+    echo 'Creating image.'
+    base_obj_type="images"
+    instance_disk_args='--image-project=my-project --image=my-image-install --boot-disk-size=40G --boot-disk-type=pd-ssd'
+    execute_with_retries       gcloud compute images create my-image-install       --project=my-project       --source-image=projects/cloud-dataproc/global/images/dataproc-1-4-deb9-20190510-000000-rc01       ${cert_args}       --storage-location=us-east1       --family=debian9
+    touch "/tmp/custom-image-my-image-20190611-160823/disk_created"
+  else
+    echo 'Creating disk.'
+    base_obj_type="disks"
+    instance_disk_args='--disk=auto-delete=yes,boot=yes,mode=rw,name=my-image-install'
+    execute_with_retries gcloud compute disks create my-image-install       --project=my-project       --zone=us-west1-a       --image=projects/cloud-dataproc/global/images/dataproc-1-4-deb9-20190510-000000-rc01       --type=pd-ssd       --size=40GB
+    touch "/tmp/custom-image-my-image-20190611-160823/disk_created"
+  fi
 
+  date
   echo 'Creating VM instance to run customization script.'
-  gcloud compute instances create my-image-install       --project=my-project       --zone=us-west1-a              --subnet=my-subnet       --no-address       --machine-type=n1-standard-2       --disk=auto-delete=yes,boot=yes,mode=rw,name=my-image-install       --accelerator=type=nvidia-tesla-v100,count=2 --maintenance-policy terminate       --service-account=my-service-account       --scopes=cloud-platform       --metadata=shutdown-timer-in-sec=500,custom-sources-path=gs://my-bucket/custom-image-my-image-20190611-160823/sources,key1=value1,key2=value2       --metadata-from-file startup-script=startup_script/run.sh
+  execute_with_retries gcloud compute instances create my-image-install       --project=my-project       --zone=us-west1-a              --subnet=my-subnet       --no-address       --machine-type=n1-standard-2       ${instance_disk_args}       --accelerator=type=nvidia-tesla-v100,count=2 --maintenance-policy terminate       --service-account=my-service-account       --scopes=cloud-platform              --metadata=shutdown-timer-in-sec=500,custom-sources-path=gs://my-bucket/custom-image-my-image-20190611-160823/sources,universe-domain=googleapis.com,dataproc-region="us-west1",key1=value1,key2=value2       --metadata-from-file startup-script=startup_script/run.sh
+
   touch /tmp/custom-image-my-image-20190611-160823/vm_created
 
+  # clean up intermediate install image
+  if [[ "${base_obj_type}" == "images" ]] ; then
+    gcloud compute images delete -q my-image-install --project=my-project
+  fi
+
+  echo "Monitor startup logs in /tmp/custom-image-my-image-20190611-160823/logs/startup-script.log"
   echo 'Waiting for customization script to finish and VM shutdown.'
-  gcloud compute instances tail-serial-port-output my-image-install       --project=my-project       --zone=us-west1-a       --port=1 2>&1       | grep 'startup-script'       | tee /tmp/custom-image-my-image-20190611-160823/logs/startup-script.log       || true
+  set -x
+  # too many serial port output requests per minute occur if they all occur at once
+  sleep $(( ( RANDOM % 60 ) + 20 ))
 
+  execute_with_retries gcloud compute instances tail-serial-port-output my-image-install       --project=my-project       --zone=us-west1-a       --port=1 2>&1       | grep 'startup-script' | grep -v '^\['       | sed -e 's/ my-image-install.*startup-script://g'       | dd bs=1 status=none of=/tmp/custom-image-my-image-20190611-160823/logs/startup-script.log       || true
   echo 'Checking customization script result.'
-  if grep 'BuildFailed:' /tmp/custom-image-my-image-20190611-160823/logs/startup-script.log; then
+  date
+  if grep -q 'BuildSucceeded:' /tmp/custom-image-my-image-20190611-160823/logs/startup-script.log; then
+    echo -e "${GREEN}Customization script succeeded.${NC}"
+  elif grep -q 'BuildFailed:' /tmp/custom-image-my-image-20190611-160823/logs/startup-script.log; then
     echo -e "${RED}Customization script failed.${NC}"
+    echo "See /tmp/custom-image-my-image-20190611-160823/logs/startup-script.log for details"
     exit 1
-  elif grep 'BuildSucceeded:' /tmp/custom-image-my-image-20190611-160823/logs/startup-script.log; then
-    echo -e "${GREEN}Customization script succeeded.${NC}"
   else
     echo 'Unable to determine the customization script result.'
     exit 1
   fi
 
+  date
   echo 'Creating custom image.'
-  gcloud compute images create my-image       --project=my-project       --source-disk-zone=us-west1-a       --source-disk=my-image-install       --storage-location=us-east1       --family=debian9
+  execute_with_retries gcloud compute images create my-image     --project=my-project     --source-disk-zone=us-west1-a     --source-disk=my-image-install     --storage-location=us-east1     --family=debian9
+
   touch /tmp/custom-image-my-image-20190611-160823/image_created
 }
 
 trap exit_handler EXIT
 mkdir -p /tmp/custom-image-my-image-20190611-160823/logs
+prepare
 main "$@" 2>&1 | tee /tmp/custom-image-my-image-20190611-160823/logs/workflow.log
 """
 
 
 class TestShellScriptGenerator(unittest.TestCase):
-  def test_generate_shell_script(self):
-    args = {
-        'run_id': 'custom-image-my-image-20190611-160823',
-        'family': 'debian9',
-        'image_name': 'my-image',
-        'customization_script': '/tmp/my-script.sh',
-        'metadata': 'key1=value1,key2=value2',
-        'extra_sources': {"ext'ra_src.txt": "/path/to/extra.txt"},
-        'machine_type': 'n1-standard-2',
-        'disk_size': 40,
-        'accelerator': 'type=nvidia-tesla-v100,count=2',
-        'gcs_bucket': 'gs://my-bucket',
-        'network': 'my-network',
-        'subnetwork': 'my-subnet',
-        'no_external_ip': True,
-        'zone': 'us-west1-a',
-        'dataproc_base_image':
-          'projects/cloud-dataproc/global/images/dataproc-1-4-deb9-20190510-000000-rc01',
-        'service_account': 'my-service-account',
-        'oauth': '',
-        'project_id': 'my-project',
-        'storage_location': 'us-east1',
-        'shutdown_timer_in_sec': 500,
-        'base_image_family': 'projects/my-dataproc-project/global/images/family/debian-10'
-    }
-
-    script = shell_script_generator.Generator().generate(args)
-
-    self.assertEqual(script, _expected_script)
-
-
-if __name__ == '__main__':
-  unittest.main()
+    def test_generate_shell_script(self):
+        args = {
+            "run_id": "custom-image-my-image-20190611-160823",
+            "family": "debian9",
+            "image_name": "my-image",
+            "customization_script": "/tmp/my-script.sh",
+            "metadata": "key1=value1,key2=value2",
+            "extra_sources": {"ext'ra_src.txt": "/path/to/extra.txt"},
+            "machine_type": "n1-standard-2",
+            "disk_size": 40,
+            "accelerator": "type=nvidia-tesla-v100,count=2",
+            "gcs_bucket": "gs://my-bucket",
+            "network": "my-network",
+            "subnetwork": "my-subnet",
+            "no_external_ip": True,
+            "zone": "us-west1-a",
+            "dataproc_base_image": "projects/cloud-dataproc/global/images/dataproc-1-4-deb9-20190510-000000-rc01",
+            "service_account": "my-service-account",
+            "oauth": "",
+            "project_id": "my-project",
+            "storage_location": "us-east1",
+            "shutdown_timer_in_sec": 500,
+            "base_image_family": "projects/my-dataproc-project/global/images/family/debian-10",
+            "trusted_cert": "",
+        }
+
+        script = shell_script_generator.Generator().generate(args)
+
+        self.assertEqual(script.strip(), _expected_script.strip())
+
+
+if __name__ == "__main__":
+    unittest.main()