diff --git a/src/google/adk/cli/cli_deploy.py b/src/google/adk/cli/cli_deploy.py index 47a5b688be..00d97d224a 100644 --- a/src/google/adk/cli/cli_deploy.py +++ b/src/google/adk/cli/cli_deploy.py @@ -17,6 +17,7 @@ import importlib import json import os +import re import shutil import subprocess import sys @@ -40,6 +41,92 @@ _AGENT_ENGINE_REQUIREMENT: Final[str] = ( 'google-cloud-aiplatform[adk,agent_engines]' ) +# Full Cloud Build private worker pool resource name, e.g. +# projects/my-project/locations/us-central1/workerPools/my-private-pool +_WORKER_POOL_RESOURCE_RE: Final[re.Pattern[str]] = re.compile( + r'^projects/[^/]+/locations/[^/]+/workerPools/[^/]+$' +) + + +def _validate_worker_pool(worker_pool: str) -> str: + """Validates a Cloud Build worker pool resource name. + + Args: + worker_pool: Full resource name of the form + `projects/{project}/locations/{location}/workerPools/{pool}`. + + Returns: + The validated worker pool resource name. + + Raises: + click.ClickException: If the resource name is empty or malformed. + """ + worker_pool = worker_pool.strip() + if not worker_pool: + raise click.ClickException('worker_pool must be a non-empty resource name.') + if not _WORKER_POOL_RESOURCE_RE.fullmatch(worker_pool): + raise click.ClickException( + 'Invalid worker_pool resource name. Expected format:' + ' projects/{project}/locations/{location}/workerPools/{pool}.' + f' Got: {worker_pool}' + ) + return worker_pool + + +def _apply_worker_pool_to_agent_config( + agent_config: dict[str, Any], + worker_pool: Optional[str], +) -> None: + """Nests worker_pool into agent_config['build_config']. + + Supports three sources, in increasing precedence: + + 1. Existing ``build_config.worker_pool`` already in ``agent_config``. + 2. Top-level ``worker_pool`` convenience key in ``.agent_engine_config.json`` + (popped so it is not forwarded as an unknown top-level field). + 3. Explicit ``worker_pool`` argument (CLI flag), which overrides both. + + The Vertex Agent Engine SDK reads Cloud Build private pools from + ``config.build_config.worker_pool`` and maps them onto + ``spec.build_spec.worker_pool``. + """ + build_config = agent_config.get('build_config') + if build_config is None: + build_config = {} + elif not isinstance(build_config, dict): + raise click.ClickException( + 'build_config in agent platform config must be a JSON object.' + ) + else: + # Copy so we do not mutate a shared structure unexpectedly. + build_config = dict(build_config) + + config_worker_pool = agent_config.pop('worker_pool', None) + if config_worker_pool is not None: + if not isinstance(config_worker_pool, str): + raise click.ClickException( + 'worker_pool in agent platform config must be a string resource name.' + ) + build_config['worker_pool'] = _validate_worker_pool(config_worker_pool) + + if worker_pool is not None: + if build_config.get('worker_pool'): + click.echo( + 'Overriding build_config.worker_pool in agent platform config with' + f' {worker_pool}' + ) + build_config['worker_pool'] = _validate_worker_pool(worker_pool) + + # Validate any worker_pool that was already nested under build_config. + if 'worker_pool' in build_config and build_config['worker_pool'] is not None: + build_config['worker_pool'] = _validate_worker_pool( + str(build_config['worker_pool']) + ) + + if build_config: + agent_config['build_config'] = build_config + else: + agent_config.pop('build_config', None) def _ensure_agent_engine_dependency(requirements_txt_path: str) -> None: @@ -886,6 +973,7 @@ def to_agent_engine( artifact_service_uri: Optional[str] = None, adk_version: Optional[str] = None, extra_packages: Optional[list[str]] = None, + worker_pool: Optional[str] = None, ) -> None: """Deploys an agent to Gemini Enterprise Agent Platform. @@ -952,6 +1040,13 @@ def to_agent_engine( used. extra_packages (list[str]): Optional. Additional local file or directory paths to stage alongside the agent and make importable in the image. + worker_pool (str): Optional. Full Cloud Build private worker pool resource + name + (`projects/{project}/locations/{location}/workerPools/{pool}`). + When set, Agent Engine builds the container image on that pool so + deploys can reach private networks / comply with org build policies. + Overrides `worker_pool` / `build_config.worker_pool` from + `.agent_engine_config.json` when both are present. """ app_name = os.path.basename(agent_folder) display_name = display_name or app_name @@ -1049,6 +1144,8 @@ def to_agent_engine( ) agent_config['description'] = description + _apply_worker_pool_to_agent_config(agent_config, worker_pool) + config_extra_packages = agent_config.pop('extra_packages', None) or [] # CLI entries resolve against the invocation dir; config-file entries # against the agent folder that declared them. diff --git a/src/google/adk/cli/cli_tools_click.py b/src/google/adk/cli/cli_tools_click.py index 12edb32993..cdcc418d66 100644 --- a/src/google/adk/cli/cli_tools_click.py +++ b/src/google/adk/cli/cli_tools_click.py @@ -2702,6 +2702,19 @@ def cli_migrate_session( " Repeatable." ), ) +@click.option( + "--worker_pool", + type=str, + default=None, + help=( + "Optional. Cloud Build private worker pool resource name used to build" + " the Agent Engine container image. Format:" + " projects/{project}/locations/{location}/workerPools/{pool}." + " Required for VPC-SC / private-network environments that cannot use" + " the default public Cloud Build pool. Overrides `worker_pool` or" + " `build_config.worker_pool` in `.agent_engine_config.json`." + ), +) @adk_services_options(default_use_local_storage=False) @click.argument( "agent", @@ -2736,6 +2749,7 @@ def cli_deploy_agent_engine( session_service_uri: str | None = None, use_local_storage: bool = False, extra_packages: tuple[str, ...] = (), + worker_pool: str | None = None, ): """Deploys an agent to Agent Engine. @@ -2749,6 +2763,12 @@ def cli_deploy_agent_engine( # With Google Cloud Project and Region adk deploy agent_engine --project=[project] --region=[region] --display_name=[app_name] my_agent + + \b + # With a private Cloud Build worker pool (VPC-SC / private network) + adk deploy agent_engine --project=[project] --region=[region] + --worker_pool=projects/[project]/locations/[region]/workerPools/[pool] + my_agent """ logging.getLogger("vertexai_genai.agentengines").setLevel(logging.INFO) try: @@ -2783,6 +2803,7 @@ def cli_deploy_agent_engine( session_service_uri=session_service_uri, adk_version=adk_version, extra_packages=list(extra_packages), + worker_pool=worker_pool, ) except Exception as e: click.secho(f"Deploy failed: {e}", fg="red", err=True) diff --git a/tests/unittests/cli/utils/test_cli_deploy.py b/tests/unittests/cli/utils/test_cli_deploy.py index 98dc493687..b533957f9f 100644 --- a/tests/unittests/cli/utils/test_cli_deploy.py +++ b/tests/unittests/cli/utils/test_cli_deploy.py @@ -1164,3 +1164,178 @@ def test_to_agent_engine_extra_packages_requirements_txt_is_not_clobbered( assert (tmp_dir / "requirements.txt").read_text() == ( "some-unrelated-package\n" ) + + +_VALID_WORKER_POOL = ( + "projects/my-gcp-project/locations/us-central1/workerPools/my-private-pool" +) + + +def test_validate_worker_pool_accepts_full_resource_name() -> None: + """A well-formed Cloud Build worker pool resource name is accepted.""" + assert cli_deploy._validate_worker_pool(_VALID_WORKER_POOL) == ( + _VALID_WORKER_POOL + ) + + +@pytest.mark.parametrize( + "bad_pool", + [ + "", + " ", + "my-private-pool", + "projects/p/locations/l/workerPools/", + "projects/p/locations/l/pools/my-pool", + "projects/p/workerPools/my-pool", + ], +) +def test_validate_worker_pool_rejects_malformed_names(bad_pool: str) -> None: + """Malformed worker pool resource names raise a clear ClickException.""" + with pytest.raises(click.ClickException): + cli_deploy._validate_worker_pool(bad_pool) + + +def test_apply_worker_pool_nests_cli_value_into_build_config() -> None: + """CLI worker_pool is nested under build_config for the Vertex SDK.""" + agent_config: Dict[str, Any] = {} + cli_deploy._apply_worker_pool_to_agent_config( + agent_config, _VALID_WORKER_POOL + ) + assert agent_config["build_config"]["worker_pool"] == _VALID_WORKER_POOL + assert "worker_pool" not in agent_config + + +def test_apply_worker_pool_pops_top_level_config_key() -> None: + """Top-level worker_pool in .agent_engine_config.json is nested and removed.""" + agent_config: Dict[str, Any] = {"worker_pool": _VALID_WORKER_POOL} + cli_deploy._apply_worker_pool_to_agent_config(agent_config, None) + assert agent_config["build_config"]["worker_pool"] == _VALID_WORKER_POOL + assert "worker_pool" not in agent_config + + +def test_apply_worker_pool_cli_overrides_config_file() -> None: + """Explicit CLI worker_pool overrides values from the config file.""" + override = "projects/other/locations/europe-west1/workerPools/compliance-pool" + agent_config: Dict[str, Any] = { + "build_config": {"worker_pool": _VALID_WORKER_POOL}, + } + cli_deploy._apply_worker_pool_to_agent_config(agent_config, override) + assert agent_config["build_config"]["worker_pool"] == override + + +def test_apply_worker_pool_preserves_other_build_config_fields() -> None: + """Existing build_config.service_account is kept when adding worker_pool.""" + agent_config: Dict[str, Any] = { + "build_config": { + "service_account": "builder@example.iam.gserviceaccount.com", + }, + } + cli_deploy._apply_worker_pool_to_agent_config( + agent_config, _VALID_WORKER_POOL + ) + assert agent_config["build_config"] == { + "service_account": "builder@example.iam.gserviceaccount.com", + "worker_pool": _VALID_WORKER_POOL, + } + + +def test_to_agent_engine_forwards_worker_pool_in_update_config( + monkeypatch: pytest.MonkeyPatch, + agent_dir: Callable[[bool, bool], Path], +) -> None: + """to_agent_engine puts worker_pool under build_config on agent_engines.update.""" + monkeypatch.setattr(shutil, "rmtree", _Recorder()) + captured: List[Dict[str, Any]] = [] + monkeypatch.setitem( + sys.modules, "vertexai", _make_recording_vertexai(captured) + ) + src_dir = agent_dir(False, False) + + cli_deploy.to_agent_engine( + agent_folder=str(src_dir), + temp_folder="tmp", + project="my-gcp-project", + region="us-central1", + adk_version="1.2.0", + worker_pool=_VALID_WORKER_POOL, + ) + + assert len(captured) == 1 + assert captured[0]["build_config"]["worker_pool"] == _VALID_WORKER_POOL + + +def test_to_agent_engine_reads_worker_pool_from_config_file( + monkeypatch: pytest.MonkeyPatch, + agent_dir: Callable[[bool, bool], Path], +) -> None: + """worker_pool from .agent_engine_config.json is forwarded on deploy.""" + monkeypatch.setattr(shutil, "rmtree", _Recorder()) + captured: List[Dict[str, Any]] = [] + monkeypatch.setitem( + sys.modules, "vertexai", _make_recording_vertexai(captured) + ) + src_dir = agent_dir(False, False) + (src_dir / ".agent_engine_config.json").write_text( + json.dumps({"worker_pool": _VALID_WORKER_POOL}) + ) + + cli_deploy.to_agent_engine( + agent_folder=str(src_dir), + temp_folder="tmp", + project="my-gcp-project", + region="us-central1", + adk_version="1.2.0", + ) + + assert captured[0]["build_config"]["worker_pool"] == _VALID_WORKER_POOL + assert "worker_pool" not in captured[0] + + +def test_to_agent_engine_rejects_invalid_worker_pool( + monkeypatch: pytest.MonkeyPatch, + agent_dir: Callable[[bool, bool], Path], +) -> None: + """An invalid --worker_pool value fails before calling Agent Engine APIs.""" + monkeypatch.setattr(shutil, "rmtree", _Recorder()) + captured: List[Dict[str, Any]] = [] + monkeypatch.setitem( + sys.modules, "vertexai", _make_recording_vertexai(captured) + ) + src_dir = agent_dir(False, False) + + with pytest.raises(click.ClickException) as exc_info: + cli_deploy.to_agent_engine( + agent_folder=str(src_dir), + temp_folder="tmp", + project="my-gcp-project", + region="us-central1", + adk_version="1.2.0", + worker_pool="not-a-resource-name", + ) + + assert "Invalid worker_pool" in str(exc_info.value) + assert captured == [] + + +def test_cli_deploy_agent_engine_passes_worker_pool(tmp_path: Path) -> None: + """--worker_pool reaches to_agent_engine as a keyword argument.""" + agent_dir = tmp_path / "my_agent" + agent_dir.mkdir() + runner = CliRunner() + with mock.patch( + "src.google.adk.cli.cli_deploy.to_agent_engine" + ) as mock_to_agent_engine: + result = runner.invoke( + cli_tools_click.main, + [ + "deploy", + "agent_engine", + f"--worker_pool={_VALID_WORKER_POOL}", + str(agent_dir), + ], + catch_exceptions=False, + ) + assert result.exit_code == 0 + mock_to_agent_engine.assert_called_once() + _, kwargs = mock_to_agent_engine.call_args + assert kwargs["worker_pool"] == _VALID_WORKER_POOL