feat: adding build image to deploy cloud_run options (#2502)

* feat: adding build image to deploy cloud_run options

Gives the ability for a user to set the build image for the deployment step to Cloud Run.  Currently it is hard coded to python:3.11-slim, and this is still the default, but this allows that value to be overriden.

* fix: applied formatting scripts

testing:

Added tests to ensure the behavior of the cli remains consistent with when used or omitted.

* chore: next time run the formatter before you commit.

---------

Co-authored-by: Ivan Cheung <ivans.mailbox@gmail.com>

src/google/adk/cli/cli_deploy.py

@@ -21,8 +21,10 @@ from typing import Optional
 import click
 from packaging.version import parse
 
+BASE_BUILD_IMAGE = 'python:3.11-slim'
+
 _DOCKERFILE_TEMPLATE = """
-FROM python:3.11-slim
+FROM {build_image}
 WORKDIR /app
 
 # Create a non-root user
@@ -130,6 +132,7 @@ def to_cloud_run(
     session_service_uri: Optional[str] = None,
     artifact_service_uri: Optional[str] = None,
     memory_service_uri: Optional[str] = None,
+    build_image: Optional[str] = BASE_BUILD_IMAGE,
     a2a: bool = False,
 ):
   """Deploys an agent to Google Cloud Run.
@@ -163,6 +166,7 @@ def to_cloud_run(
     session_service_uri: The URI of the session service.
     artifact_service_uri: The URI of the artifact service.
     memory_service_uri: The URI of the memory service.
+    build_image: The image to use for building the Dockerfile.
   """
   app_name = app_name or os.path.basename(agent_folder)
 
@@ -211,6 +215,7 @@ def to_cloud_run(
         adk_version=adk_version,
         host_option=host_option,
         a2a_option=a2a_option,
+        build_image=build_image,
     )
     dockerfile_path = os.path.join(temp_folder, 'Dockerfile')
     os.makedirs(temp_folder, exist_ok=True)
@@ -474,6 +479,7 @@ def to_gke(
     session_service_uri: Optional[str] = None,
     artifact_service_uri: Optional[str] = None,
     memory_service_uri: Optional[str] = None,
+    build_image: Optional[str] = BASE_BUILD_IMAGE,
     a2a: bool = False,
 ):
   """Deploys an agent to Google Kubernetes Engine(GKE).
@@ -495,6 +501,7 @@ def to_gke(
     session_service_uri: The URI of the session service.
     artifact_service_uri: The URI of the artifact service.
     memory_service_uri: The URI of the memory service.
+    build_image: The image to use for building the Dockerfile.
   """
   click.secho(
       '\n🚀 Starting ADK Agent Deployment to GKE...', fg='cyan', bold=True
@@ -556,6 +563,7 @@ def to_gke(
         adk_version=adk_version,
         host_option=host_option,
         a2a_option='--a2a' if a2a else '',
+        build_image=build_image,
     )
     dockerfile_path = os.path.join(temp_folder, 'Dockerfile')
     os.makedirs(temp_folder, exist_ok=True)

src/google/adk/cli/cli_tools_click.py

@@ -44,6 +44,8 @@ LOG_LEVELS = click.Choice(
     case_sensitive=False,
 )
 
+BASE_BUILD_IMAGE = "python:3.11-slim"
+
 
 class HelpfulCommand(click.Command):
   """Command that shows full help on error instead of just the error message.
@@ -968,6 +970,13 @@ def cli_api_server(
     help="Optional. Any additional origins to allow for CORS.",
     multiple=True,
 )
+@click.option(
+    "--build_image",
+    type=str,
+    default=BASE_BUILD_IMAGE,
+    show_default=True,
+    help="Optional. The docker build version used in Cloud Run deployment. ",
+)
 # TODO: Add eval_storage_uri option back when evals are supported in Cloud Run.
 @adk_services_options()
 @deprecated_adk_services_options()
@@ -991,6 +1000,7 @@ def cli_deploy_cloud_run(
     session_db_url: Optional[str] = None,  # Deprecated
     artifact_storage_uri: Optional[str] = None,  # Deprecated
     a2a: bool = False,
+    build_image: Optional[str] = BASE_BUILD_IMAGE,
 ):
   """Deploys an agent to Cloud Run.
 
@@ -1029,6 +1039,7 @@ def cli_deploy_cloud_run(
         artifact_service_uri=artifact_service_uri,
         memory_service_uri=memory_service_uri,
         a2a=a2a,
+        build_image=build_image,
     )
   except Exception as e:
     click.secho(f"Deploy failed: {e}", fg="red", err=True)
@@ -1281,6 +1292,13 @@ def cli_deploy_agent_engine(
         " version in the dev environment)"
     ),
 )
+@click.option(
+    "--build_image",
+    type=str,
+    default=BASE_BUILD_IMAGE,
+    show_default=True,
+    help="Optional. The docker build version used in GKE deployment. ",
+)
 @adk_services_options()
 @click.argument(
     "agent",
@@ -1304,6 +1322,7 @@ def cli_deploy_gke(
     session_service_uri: Optional[str] = None,
     artifact_service_uri: Optional[str] = None,
     memory_service_uri: Optional[str] = None,
+    build_image: Optional[str] = BASE_BUILD_IMAGE,
 ):
   """Deploys an agent to GKE.
 
@@ -1330,6 +1349,7 @@ def cli_deploy_gke(
         session_service_uri=session_service_uri,
         artifact_service_uri=artifact_service_uri,
         memory_service_uri=memory_service_uri,
+        build_image=build_image,
     )
   except Exception as e:
     click.secho(f"Deploy failed: {e}", fg="red", err=True)

tests/unittests/cli/utils/test_cli_tools_click.py

@@ -168,11 +168,43 @@ def test_cli_deploy_cloud_run_success(
           "proj",
           "--region",
           "asia-northeast1",
+          "--build_image",
+          "my-custom-image",
           str(agent_dir),
       ],
   )
   assert result.exit_code == 0
   assert rec.calls, "cli_deploy.to_cloud_run must be invoked"
+  called_kwargs = rec.calls[0][1]
+  assert called_kwargs.get("build_image") == "my-custom-image"
+
+
+def test_cli_deploy_cloud_run_default_build_image(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+  """Test that the default build_image is used."""
+  rec = _Recorder()
+  monkeypatch.setattr(cli_tools_click.cli_deploy, "to_cloud_run", rec)
+
+  agent_dir = tmp_path / "agent2"
+  agent_dir.mkdir()
+  runner = CliRunner()
+  result = runner.invoke(
+      cli_tools_click.main,
+      [
+          "deploy",
+          "cloud_run",
+          "--project",
+          "proj",
+          "--region",
+          "asia-northeast1",
+          str(agent_dir),
+      ],
+  )
+  assert result.exit_code == 0
+  assert rec.calls, "cli_deploy.to_cloud_run must be invoked"
+  called_kwargs = rec.calls[0][1]
+  assert called_kwargs.get("build_image") == cli_tools_click.BASE_BUILD_IMAGE
 
 
 def test_cli_deploy_cloud_run_failure(
@@ -251,6 +283,8 @@ def test_cli_deploy_gke_success(
           "us-central1",
           "--cluster_name",
           "my-cluster",
+          "--build_image",
+          "my-gke-image",
           str(agent_dir),
       ],
   )
@@ -260,6 +294,37 @@ def test_cli_deploy_gke_success(
   assert called_kwargs.get("project") == "test-proj"
   assert called_kwargs.get("region") == "us-central1"
   assert called_kwargs.get("cluster_name") == "my-cluster"
+  assert called_kwargs.get("build_image") == "my-gke-image"
+
+
+def test_cli_deploy_gke_default_build_image(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+  """Test that the default build_image is used for gke."""
+  rec = _Recorder()
+  monkeypatch.setattr(cli_tools_click.cli_deploy, "to_gke", rec)
+
+  agent_dir = tmp_path / "agent_gke"
+  agent_dir.mkdir()
+  runner = CliRunner()
+  result = runner.invoke(
+      cli_tools_click.main,
+      [
+          "deploy",
+          "gke",
+          "--project",
+          "test-proj",
+          "--region",
+          "us-central1",
+          "--cluster_name",
+          "my-cluster",
+          str(agent_dir),
+      ],
+  )
+  assert result.exit_code == 0
+  assert rec.calls, "cli_deploy.to_gke must be invoked"
+  called_kwargs = rec.calls[0][1]
+  assert called_kwargs.get("build_image") == cli_tools_click.BASE_BUILD_IMAGE
 
 
 # cli eval