feat: cli funcionality to deploy an Agent to a running GKE cluster

Merge https://github.com/google/adk-python/pull/1607

- Added CLI functionality so that we can deploy and Agent onto a GKE cluster
- Related documentation https://github.com/google/adk-docs/pull/445

COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/1607 from vicentefb:GkeDeployAgent 42f35d93b0a5df5f6dbeb9ee7f869cde51e2f6eb
PiperOrigin-RevId: 786857789
This commit is contained in:
Vicente Ferrara
2025-07-24 15:32:23 -07:00
committed by Copybara-Service
parent c8f8b4a20a
commit a858d79b3a
4 changed files with 1297 additions and 353 deletions
+231 -3
View File
@@ -153,11 +153,11 @@ def to_cloud_run(
app_name: The name of the app, by default, it's basename of `agent_folder`.
temp_folder: The temp folder for the generated Cloud Run source files.
port: The port of the ADK api server.
allow_origins: The list of allowed origins for the ADK api server.
trace_to_cloud: Whether to enable Cloud Trace.
with_ui: Whether to deploy with UI.
verbosity: The verbosity level of the CLI.
adk_version: The ADK version to use in Cloud Run.
allow_origins: The list of allowed origins for the ADK api server.
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.
@@ -182,7 +182,7 @@ def to_cloud_run(
if os.path.exists(requirements_txt_path)
else ''
)
click.echo('Copying agent source code complete.')
click.echo('Copying agent source code completed.')
# create Dockerfile
click.echo('Creating Dockerfile...')
@@ -425,7 +425,7 @@ def to_agent_engine(
'async_stream': ['async_stream_query'],
'stream': ['stream_query', 'streaming_agent_run_with_events'],
},
sys_paths=[temp_folder[1:]],
sys_paths=[temp_folder],
)
agent_config = dict(
agent_engine=agent_engine,
@@ -443,3 +443,231 @@ def to_agent_engine(
finally:
click.echo(f'Cleaning up the temp folder: {temp_folder}')
shutil.rmtree(temp_folder)
def to_gke(
*,
agent_folder: str,
project: Optional[str],
region: Optional[str],
cluster_name: str,
service_name: str,
app_name: str,
temp_folder: str,
port: int,
trace_to_cloud: bool,
with_ui: bool,
log_level: str,
verbosity: str,
adk_version: str,
allow_origins: Optional[list[str]] = None,
session_service_uri: Optional[str] = None,
artifact_service_uri: Optional[str] = None,
memory_service_uri: Optional[str] = None,
a2a: bool = False,
):
"""Deploys an agent to Google Kubernetes Engine(GKE).
Args:
agent_folder: The folder (absolute path) containing the agent source code.
project: Google Cloud project id.
region: Google Cloud region.
cluster_name: The name of the GKE cluster.
service_name: The service name in GKE.
app_name: The name of the app, by default, it's basename of `agent_folder`.
temp_folder: The local directory to use as a temporary workspace for preparing deployment artifacts. The tool populates this folder with a copy of the agent's source code and auto-generates necessary files like a Dockerfile and deployment.yaml.
port: The port of the ADK api server.
trace_to_cloud: Whether to enable Cloud Trace.
with_ui: Whether to deploy with UI.
verbosity: The verbosity level of the CLI.
adk_version: The ADK version to use in GKE.
allow_origins: The list of allowed origins for the ADK api server.
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.
"""
click.secho(
'\n🚀 Starting ADK Agent Deployment to GKE...', fg='cyan', bold=True
)
click.echo('--------------------------------------------------')
# Resolve project early to show the user which one is being used
project = _resolve_project(project)
click.echo(f' Project: {project}')
click.echo(f' Region: {region}')
click.echo(f' Cluster: {cluster_name}')
click.echo('--------------------------------------------------\n')
app_name = app_name or os.path.basename(agent_folder)
click.secho('STEP 1: Preparing build environment...', bold=True)
click.echo(f' - Using temporary directory: {temp_folder}')
# remove temp_folder if exists
if os.path.exists(temp_folder):
click.echo(' - Removing existing temporary directory...')
shutil.rmtree(temp_folder)
try:
# copy agent source code
click.echo(' - Copying agent source code...')
agent_src_path = os.path.join(temp_folder, 'agents', app_name)
shutil.copytree(agent_folder, agent_src_path)
requirements_txt_path = os.path.join(agent_src_path, 'requirements.txt')
install_agent_deps = (
f'RUN pip install -r "/app/agents/{app_name}/requirements.txt"'
if os.path.exists(requirements_txt_path)
else ''
)
click.secho('✅ Environment prepared.', fg='green')
allow_origins_option = (
f'--allow_origins={",".join(allow_origins)}' if allow_origins else ''
)
# create Dockerfile
click.secho('\nSTEP 2: Generating deployment files...', bold=True)
click.echo(' - Creating Dockerfile...')
host_option = '--host=0.0.0.0' if adk_version > '0.5.0' else ''
dockerfile_content = _DOCKERFILE_TEMPLATE.format(
gcp_project_id=project,
gcp_region=region,
app_name=app_name,
port=port,
command='web' if with_ui else 'api_server',
install_agent_deps=install_agent_deps,
service_option=_get_service_option_by_adk_version(
adk_version,
session_service_uri,
artifact_service_uri,
memory_service_uri,
),
trace_to_cloud_option='--trace_to_cloud' if trace_to_cloud else '',
allow_origins_option=allow_origins_option,
adk_version=adk_version,
host_option=host_option,
a2a_option='--a2a' if a2a else '',
)
dockerfile_path = os.path.join(temp_folder, 'Dockerfile')
os.makedirs(temp_folder, exist_ok=True)
with open(dockerfile_path, 'w', encoding='utf-8') as f:
f.write(
dockerfile_content,
)
click.secho(f'✅ Dockerfile generated: {dockerfile_path}', fg='green')
# Build and push the Docker image
click.secho(
'\nSTEP 3: Building container image with Cloud Build...', bold=True
)
click.echo(
' (This may take a few minutes. Raw logs from gcloud will be shown'
' below.)'
)
project = _resolve_project(project)
image_name = f'gcr.io/{project}/{service_name}'
subprocess.run(
[
'gcloud',
'builds',
'submit',
'--tag',
image_name,
'--verbosity',
log_level.lower() if log_level else verbosity,
temp_folder,
],
check=True,
)
click.secho('✅ Container image built and pushed successfully.', fg='green')
# Create a Kubernetes deployment
click.echo(' - Creating Kubernetes deployment.yaml...')
deployment_yaml = f"""
apiVersion: apps/v1
kind: Deployment
metadata:
name: {service_name}
labels:
app.kubernetes.io/name: adk-agent
app.kubernetes.io/version: {adk_version}
app.kubernetes.io/instance: {service_name}
app.kubernetes.io/managed-by: adk-cli
spec:
replicas: 1
selector:
matchLabels:
app: {service_name}
template:
metadata:
labels:
app: {service_name}
app.kubernetes.io/name: adk-agent
app.kubernetes.io/version: {adk_version}
app.kubernetes.io/instance: {service_name}
app.kubernetes.io/managed-by: adk-cli
spec:
containers:
- name: {service_name}
image: {image_name}
ports:
- containerPort: {port}
---
apiVersion: v1
kind: Service
metadata:
name: {service_name}
spec:
type: LoadBalancer
selector:
app: {service_name}
ports:
- port: 80
targetPort: {port}
"""
deployment_yaml_path = os.path.join(temp_folder, 'deployment.yaml')
with open(deployment_yaml_path, 'w', encoding='utf-8') as f:
f.write(deployment_yaml)
click.secho(
f'✅ Kubernetes deployment manifest generated: {deployment_yaml_path}',
fg='green',
)
# Apply the deployment
click.secho('\nSTEP 4: Applying deployment to GKE cluster...', bold=True)
click.echo(' - Getting cluster credentials...')
subprocess.run(
[
'gcloud',
'container',
'clusters',
'get-credentials',
cluster_name,
'--region',
region,
'--project',
project,
],
check=True,
)
click.echo(' - Applying Kubernetes manifest...')
result = subprocess.run(
['kubectl', 'apply', '-f', temp_folder],
check=True,
capture_output=True, # <-- Add this
text=True, # <-- Add this
)
# 2. Print the captured output line by line
click.secho(
' - The following resources were applied to the cluster:', fg='green'
)
for line in result.stdout.strip().split('\n'):
click.echo(f' - {line}')
finally:
click.secho('\nSTEP 5: Cleaning up...', bold=True)
click.echo(f' - Removing temporary directory: {temp_folder}')
shutil.rmtree(temp_folder)
click.secho(
'\n🎉 Deployment to GKE finished successfully!', fg='cyan', bold=True
)
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
+150 -109
View File
@@ -23,44 +23,16 @@ from types import SimpleNamespace
from typing import Any
from typing import Dict
from typing import List
from typing import Optional
from typing import Tuple
from unittest import mock
import click
from click.testing import CliRunner
from google.adk.agents.base_agent import BaseAgent
from google.adk.cli import cli_tools_click
from google.adk.evaluation.eval_case import EvalCase
from google.adk.evaluation.eval_set import EvalSet
from google.adk.evaluation.local_eval_set_results_manager import LocalEvalSetResultsManager
from google.adk.evaluation.local_eval_sets_manager import LocalEvalSetsManager
import google.adk.evaluation.local_eval_sets_manager as managerModule
from pydantic import BaseModel
import pytest
class DummyAgent(BaseAgent):
def __init__(self, name):
super().__init__(name=name)
self.sub_agents = []
root_agent = DummyAgent(name="dummy_agent")
@pytest.fixture
def mock_load_eval_set_from_file():
with mock.patch(
"google.adk.evaluation.local_eval_sets_manager.load_eval_set_from_file"
) as mock_func:
yield mock_func
@pytest.fixture
def mock_get_root_agent():
with mock.patch("google.adk.cli.cli_eval.get_root_agent") as mock_func:
mock_func.return_value = root_agent
yield mock_func
from src.google.adk.cli import cli_tools_click
# Helpers
@@ -78,13 +50,14 @@ class _Recorder(BaseModel):
def _mute_click(monkeypatch: pytest.MonkeyPatch) -> None:
"""Suppress click output during tests."""
monkeypatch.setattr(click, "echo", lambda *a, **k: None)
monkeypatch.setattr(click, "secho", lambda *a, **k: None)
# Keep secho for error messages
# monkeypatch.setattr(click, "secho", lambda *a, **k: None)
# validate_exclusive
def test_validate_exclusive_allows_single() -> None:
"""Providing exactly one exclusive option should pass."""
ctx = click.Context(cli_tools_click.main)
ctx = click.Context(cli_tools_click.cli_run)
param = SimpleNamespace(name="replay")
assert (
cli_tools_click.validate_exclusive(ctx, param, "file.json") == "file.json"
@@ -93,7 +66,7 @@ def test_validate_exclusive_allows_single() -> None:
def test_validate_exclusive_blocks_multiple() -> None:
"""Providing two exclusive options should raise UsageError."""
ctx = click.Context(cli_tools_click.main)
ctx = click.Context(cli_tools_click.cli_run)
param1 = SimpleNamespace(name="replay")
param2 = SimpleNamespace(name="resume")
@@ -184,10 +157,6 @@ def test_cli_deploy_cloud_run_failure(
monkeypatch.setattr(cli_tools_click.cli_deploy, "to_cloud_run", _boom)
# intercept click.secho(error=True) output
captured: List[str] = []
monkeypatch.setattr(click, "secho", lambda msg, **__: captured.append(msg))
agent_dir = tmp_path / "agent3"
agent_dir.mkdir()
runner = CliRunner()
@@ -196,7 +165,73 @@ def test_cli_deploy_cloud_run_failure(
)
assert result.exit_code == 0
assert any("Deploy failed: boom" in m for m in captured)
assert "Deploy failed: boom" in result.output
# cli deploy agent_engine
def test_cli_deploy_agent_engine_success(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
"""Successful path should call cli_deploy.to_agent_engine."""
rec = _Recorder()
monkeypatch.setattr(cli_tools_click.cli_deploy, "to_agent_engine", rec)
agent_dir = tmp_path / "agent_ae"
agent_dir.mkdir()
runner = CliRunner()
result = runner.invoke(
cli_tools_click.main,
[
"deploy",
"agent_engine",
"--project",
"test-proj",
"--region",
"us-central1",
"--staging_bucket",
"gs://mybucket",
str(agent_dir),
],
)
assert result.exit_code == 0
assert rec.calls, "cli_deploy.to_agent_engine must be invoked"
called_kwargs = rec.calls[0][1]
assert called_kwargs.get("project") == "test-proj"
assert called_kwargs.get("region") == "us-central1"
assert called_kwargs.get("staging_bucket") == "gs://mybucket"
# cli deploy gke
def test_cli_deploy_gke_success(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
"""Successful path should call cli_deploy.to_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("project") == "test-proj"
assert called_kwargs.get("region") == "us-central1"
assert called_kwargs.get("cluster_name") == "my-cluster"
# cli eval
@@ -204,16 +239,30 @@ def test_cli_eval_missing_deps_raises(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
"""If cli_eval sub-module is missing, command should raise ClickException."""
# Ensure .cli_eval is not importable
orig_import = builtins.__import__
def _fake_import(name: str, *a: Any, **k: Any):
if name.endswith(".cli_eval") or name == "google.adk.cli.cli_eval":
raise ModuleNotFoundError()
return orig_import(name, *a, **k)
def _fake_import(name: str, globals=None, locals=None, fromlist=(), level=0):
if name == "google.adk.cli.cli_eval" or (level > 0 and "cli_eval" in name):
raise ModuleNotFoundError(f"Simulating missing {name}")
return orig_import(name, globals, locals, fromlist, level)
monkeypatch.setattr(builtins, "__import__", _fake_import)
agent_dir = tmp_path / "agent_missing_deps"
agent_dir.mkdir()
(agent_dir / "__init__.py").touch()
eval_file = tmp_path / "dummy.json"
eval_file.touch()
runner = CliRunner()
result = runner.invoke(
cli_tools_click.main,
["eval", str(agent_dir), str(eval_file)],
)
assert result.exit_code != 0
assert isinstance(result.exception, SystemExit)
assert cli_tools_click.MISSING_EVAL_DEPENDENCIES_MESSAGE in result.output
# cli web & api_server (uvicorn patched)
@pytest.fixture()
@@ -235,18 +284,18 @@ def _patch_uvicorn(monkeypatch: pytest.MonkeyPatch) -> _Recorder:
monkeypatch.setattr(
cli_tools_click.uvicorn, "Server", lambda *_a, **_k: _DummyServer()
)
monkeypatch.setattr(
cli_tools_click, "get_fast_api_app", lambda **_k: object()
)
return rec
def test_cli_web_invokes_uvicorn(
tmp_path: Path, _patch_uvicorn: _Recorder
tmp_path: Path, _patch_uvicorn: _Recorder, monkeypatch: pytest.MonkeyPatch
) -> None:
"""`adk web` should configure and start uvicorn.Server.run."""
agents_dir = tmp_path / "agents"
agents_dir.mkdir()
monkeypatch.setattr(
cli_tools_click, "get_fast_api_app", lambda **_k: object()
)
runner = CliRunner()
result = runner.invoke(cli_tools_click.main, ["web", str(agents_dir)])
assert result.exit_code == 0
@@ -254,84 +303,76 @@ def test_cli_web_invokes_uvicorn(
def test_cli_api_server_invokes_uvicorn(
tmp_path: Path, _patch_uvicorn: _Recorder
tmp_path: Path, _patch_uvicorn: _Recorder, monkeypatch: pytest.MonkeyPatch
) -> None:
"""`adk api_server` should configure and start uvicorn.Server.run."""
agents_dir = tmp_path / "agents_api"
agents_dir.mkdir()
monkeypatch.setattr(
cli_tools_click, "get_fast_api_app", lambda **_k: object()
)
runner = CliRunner()
result = runner.invoke(cli_tools_click.main, ["api_server", str(agents_dir)])
assert result.exit_code == 0
assert _patch_uvicorn.calls, "uvicorn.Server.run must be called"
def test_cli_eval_with_eval_set_file_path(
mock_load_eval_set_from_file,
mock_get_root_agent,
tmp_path,
):
agent_path = tmp_path / "my_agent"
agent_path.mkdir()
(agent_path / "__init__.py").touch()
def test_cli_web_passes_service_uris(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, _patch_uvicorn: _Recorder
) -> None:
"""`adk web` should pass service URIs to get_fast_api_app."""
agents_dir = tmp_path / "agents"
agents_dir.mkdir()
eval_set_file = tmp_path / "my_evals.json"
eval_set_file.write_text("{}")
mock_get_app = _Recorder()
monkeypatch.setattr(cli_tools_click, "get_fast_api_app", mock_get_app)
mock_load_eval_set_from_file.return_value = EvalSet(
eval_set_id="my_evals",
eval_cases=[EvalCase(eval_id="case1", conversation=[])],
runner = CliRunner()
result = runner.invoke(
cli_tools_click.main,
[
"web",
str(agents_dir),
"--session_service_uri",
"sqlite:///test.db",
"--artifact_service_uri",
"gs://mybucket",
"--memory_service_uri",
"rag://mycorpus",
],
)
result = CliRunner().invoke(
cli_tools_click.cli_eval,
[str(agent_path), str(eval_set_file)],
)
assert result.exit_code == 0
# Assert that we wrote eval set results
eval_set_results_manager = LocalEvalSetResultsManager(
agents_dir=str(tmp_path)
)
eval_set_results = eval_set_results_manager.list_eval_set_results(
app_name="my_agent"
)
assert len(eval_set_results) == 1
assert mock_get_app.calls
called_kwargs = mock_get_app.calls[0][1]
assert called_kwargs.get("session_service_uri") == "sqlite:///test.db"
assert called_kwargs.get("artifact_service_uri") == "gs://mybucket"
assert called_kwargs.get("memory_service_uri") == "rag://mycorpus"
def test_cli_eval_with_eval_set_id(
mock_get_root_agent,
tmp_path,
):
app_name = "test_app"
eval_set_id = "test_eval_set_id"
agent_path = tmp_path / app_name
agent_path.mkdir()
(agent_path / "__init__.py").touch()
def test_cli_web_passes_deprecated_uris(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, _patch_uvicorn: _Recorder
) -> None:
"""`adk web` should use deprecated URIs if new ones are not provided."""
agents_dir = tmp_path / "agents"
agents_dir.mkdir()
eval_sets_manager = LocalEvalSetsManager(agents_dir=str(tmp_path))
eval_sets_manager.create_eval_set(app_name=app_name, eval_set_id=eval_set_id)
eval_sets_manager.add_eval_case(
app_name=app_name,
eval_set_id=eval_set_id,
eval_case=EvalCase(eval_id="case1", conversation=[]),
)
eval_sets_manager.add_eval_case(
app_name=app_name,
eval_set_id=eval_set_id,
eval_case=EvalCase(eval_id="case2", conversation=[]),
)
mock_get_app = _Recorder()
monkeypatch.setattr(cli_tools_click, "get_fast_api_app", mock_get_app)
result = CliRunner().invoke(
cli_tools_click.cli_eval,
[str(agent_path), "test_eval_set_id:case1,case2"],
runner = CliRunner()
result = runner.invoke(
cli_tools_click.main,
[
"web",
str(agents_dir),
"--session_db_url",
"sqlite:///deprecated.db",
"--artifact_storage_uri",
"gs://deprecated",
],
)
assert result.exit_code == 0
# Assert that we wrote eval set results
eval_set_results_manager = LocalEvalSetResultsManager(
agents_dir=str(tmp_path)
)
eval_set_results = eval_set_results_manager.list_eval_set_results(
app_name=app_name
)
assert len(eval_set_results) == 2
assert mock_get_app.calls
called_kwargs = mock_get_app.calls[0][1]
assert called_kwargs.get("session_service_uri") == "sqlite:///deprecated.db"
assert called_kwargs.get("artifact_service_uri") == "gs://deprecated"