Skip to content

feat: add build configuration support to Vertex AI Reasoning Engine #7046

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat: add build configuration support to Vertex AI Reasoning Engine #7046

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 22 additions & 0 deletions agentplatform/_genai/agent_engines.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1988,6 +1988,7 @@ def create(
agent_config_source=agent_config_source,
container_spec=config.container_spec,
keep_alive_probe=keep_alive_probe,
build_config=config.build_config,
)
operation = self._create(config=api_config)
reasoning_engine_id = _agent_engines_utils._get_reasoning_engine_id(
Expand Down Expand Up @@ -2301,6 +2302,7 @@ def _create_config(
container_spec: Optional[types.ReasoningEngineSpecContainerSpecDict] = None,
keep_alive_probe: Optional[dict[str, Any]] = None,
traffic_config: Optional[types.ReasoningEngineTrafficConfigDict] = None,
build_config: Optional[types.ReasoningEngineSpecBuildSpecDict] = None,
) -> types.UpdateAgentEngineConfigDict:
import sys

Expand Down Expand Up @@ -2507,6 +2509,25 @@ def _create_config(
agent_engine_spec["service_account"] = service_account
update_masks.append("spec.service_account")

if build_config is not None:
if agent_engine_spec is None:
agent_engine_spec = {}
build_spec: dict[str, Any] = {}
if isinstance(build_config, dict):
worker_pool = build_config.get("worker_pool")
build_service_account = build_config.get("service_account")
else:
worker_pool = getattr(build_config, "worker_pool", None)
build_service_account = getattr(build_config, "service_account", None)
if worker_pool is not None:
build_spec["worker_pool"] = worker_pool
update_masks.append("spec.build_spec.worker_pool")
if build_service_account is not None:
build_spec["service_account"] = build_service_account
update_masks.append("spec.build_spec.service_account")
if build_spec:
agent_engine_spec["build_spec"] = build_spec

if agent_engine_spec is not None:
config["spec"] = agent_engine_spec

Expand Down Expand Up @@ -2778,6 +2799,7 @@ def update(
container_spec=container_spec,
keep_alive_probe=keep_alive_probe,
traffic_config=traffic_config,
build_config=config.build_config,
)
operation = self._update(name=name, config=api_config)
reasoning_engine_id = _agent_engines_utils._get_reasoning_engine_id(
Expand Down
6 changes: 6 additions & 0 deletions agentplatform/_genai/types/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1462,6 +1462,9 @@
from .common import ReasoningEngineRuntimeRevisionDict
from .common import ReasoningEngineRuntimeRevisionOrDict
from .common import ReasoningEngineSpec
from .common import ReasoningEngineSpecBuildSpec
from .common import ReasoningEngineSpecBuildSpecDict
from .common import ReasoningEngineSpecBuildSpecOrDict
from .common import ReasoningEngineSpecContainerSpec
from .common import ReasoningEngineSpecContainerSpecDict
from .common import ReasoningEngineSpecContainerSpecOrDict
Expand Down Expand Up @@ -2609,6 +2612,9 @@
"ReasoningEngineSpecContainerSpec",
"ReasoningEngineSpecContainerSpecDict",
"ReasoningEngineSpecContainerSpecOrDict",
"ReasoningEngineSpecBuildSpec",
"ReasoningEngineSpecBuildSpecDict",
"ReasoningEngineSpecBuildSpecOrDict",
"ReasoningEngineSpec",
"ReasoningEngineSpecDict",
"ReasoningEngineSpecOrDict",
Expand Down
42 changes: 42 additions & 0 deletions agentplatform/_genai/types/common.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8852,6 +8852,34 @@ class ReasoningEngineSpecContainerSpecDict(TypedDict, total=False):
]


class ReasoningEngineSpecBuildSpec(_common.BaseModel):
"""Specification for building container image."""

service_account: Optional[str] = Field(
default=None,
description="""Optional. The service account that Cloud Build uses to run the build. This field is only applicable when `worker_pool` is specified (i.e., for custom worker pools). If `worker_pool` is not specified, this field is ignored and the build runs using the Google-managed service agent.""",
)
worker_pool: Optional[str] = Field(
default=None,
description="""Optional. The resource name of the Cloud Build WorkerPool to use for the build. Format: `projects/{project}/locations/{location}/workerPools/{worker_pool}`""",
)


class ReasoningEngineSpecBuildSpecDict(TypedDict, total=False):
"""Specification for building container image."""

service_account: Optional[str]
"""Optional. The service account that Cloud Build uses to run the build. This field is only applicable when `worker_pool` is specified (i.e., for custom worker pools). If `worker_pool` is not specified, this field is ignored and the build runs using the Google-managed service agent."""

worker_pool: Optional[str]
"""Optional. The resource name of the Cloud Build WorkerPool to use for the build. Format: `projects/{project}/locations/{location}/workerPools/{worker_pool}`"""


ReasoningEngineSpecBuildSpecOrDict = Union[
ReasoningEngineSpecBuildSpec, ReasoningEngineSpecBuildSpecDict
]


class ReasoningEngineSpec(_common.BaseModel):
"""The specification of an agent engine."""

Expand Down Expand Up @@ -8895,6 +8923,10 @@ class ReasoningEngineSpec(_common.BaseModel):
default=None,
description="""Deploy from a container image with a defined entrypoint and commands.""",
)
build_spec: Optional[ReasoningEngineSpecBuildSpec] = Field(
default=None,
description="""Optional. Configuration for building container image.""",
)


class ReasoningEngineSpecDict(TypedDict, total=False):
Expand Down Expand Up @@ -8930,6 +8962,9 @@ class ReasoningEngineSpecDict(TypedDict, total=False):
container_spec: Optional[ReasoningEngineSpecContainerSpecDict]
"""Deploy from a container image with a defined entrypoint and commands."""

build_spec: Optional[ReasoningEngineSpecBuildSpecDict]
"""Optional. Configuration for building container image."""


ReasoningEngineSpecOrDict = Union[ReasoningEngineSpec, ReasoningEngineSpecDict]

Expand Down Expand Up @@ -26547,6 +26582,10 @@ class AgentEngineConfig(_common.BaseModel):
traffic_config: Optional[ReasoningEngineTrafficConfig] = Field(
default=None, description="""The traffic config for the Agent Engine."""
)
build_config: Optional[ReasoningEngineSpecBuildSpec] = Field(
default=None,
description="""The build config for the Agent Engine. Allows bringing your own Cloud Build private worker pool (BYOBP) and, optionally, a build-time service account for the container build. Supported keys: `worker_pool` (the resource name of the Cloud Build WorkerPool to use for the build) and `service_account` (the service account that Cloud Build uses to run the build; only applicable when `worker_pool` is specified).""",
)


class AgentEngineConfigDict(TypedDict, total=False):
Expand Down Expand Up @@ -26734,6 +26773,9 @@ class AgentEngineConfigDict(TypedDict, total=False):
traffic_config: Optional[ReasoningEngineTrafficConfigDict]
"""The traffic config for the Agent Engine."""

build_config: Optional[ReasoningEngineSpecBuildSpecDict]
"""The build config for the Agent Engine. Allows bringing your own Cloud Build private worker pool (BYOBP) and, optionally, a build-time service account for the container build. Supported keys: `worker_pool` (the resource name of the Cloud Build WorkerPool to use for the build) and `service_account` (the service account that Cloud Build uses to run the build; only applicable when `worker_pool` is specified)."""


AgentEngineConfigOrDict = Union[AgentEngineConfig, AgentEngineConfigDict]

Expand Down
60 changes: 60 additions & 0 deletions tests/unit/agentplatform/genai/test_agent_engines.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -554,6 +554,16 @@ def register_operations(self) -> Dict[str, List[str]]:
_genai_types.IdentityType.SERVICE_ACCOUNT
)
_TEST_AGENT_ENGINE_ENCRYPTION_SPEC = {"kms_key_name": "test-kms-key"}
_TEST_AGENT_ENGINE_BUILD_WORKER_POOL = (
"projects/test-project/locations/us-central1/workerPools/test-pool"
)
_TEST_AGENT_ENGINE_BUILD_SERVICE_ACCOUNT = (
"test-build-sa@test-project.iam.gserviceaccount.com"
)
_TEST_AGENT_ENGINE_BUILD_CONFIG = _genai_types.ReasoningEngineSpecBuildSpecDict(
worker_pool=_TEST_AGENT_ENGINE_BUILD_WORKER_POOL,
service_account=_TEST_AGENT_ENGINE_BUILD_SERVICE_ACCOUNT,
)
_TEST_AGENT_ENGINE_KEEP_ALIVE_PROBE = {
"http_get": {
"path": "/health",
Expand Down Expand Up @@ -1046,6 +1056,51 @@ def test_create_agent_engine_config_full(self, mock_prepare):
== _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT
)

@mock.patch.object(_agent_engines_utils, "_prepare")
def test_create_agent_engine_config_with_build_config(self, mock_prepare):
config = self.client.agent_engines._create_config(
mode="create",
agent=self.test_agent,
staging_bucket=_TEST_STAGING_BUCKET,
requirements=_TEST_AGENT_ENGINE_REQUIREMENTS,
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
build_config=_TEST_AGENT_ENGINE_BUILD_CONFIG,
)
assert config["spec"]["build_spec"] == {
"worker_pool": _TEST_AGENT_ENGINE_BUILD_WORKER_POOL,
"service_account": _TEST_AGENT_ENGINE_BUILD_SERVICE_ACCOUNT,
}

@mock.patch.object(_agent_engines_utils, "_prepare")
def test_create_agent_engine_config_with_build_config_worker_pool_only(
self, mock_prepare
):
config = self.client.agent_engines._create_config(
mode="create",
agent=self.test_agent,
staging_bucket=_TEST_STAGING_BUCKET,
requirements=_TEST_AGENT_ENGINE_REQUIREMENTS,
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
build_config={"worker_pool": _TEST_AGENT_ENGINE_BUILD_WORKER_POOL},
)
assert config["spec"]["build_spec"] == {
"worker_pool": _TEST_AGENT_ENGINE_BUILD_WORKER_POOL,
}

@mock.patch.object(_agent_engines_utils, "_prepare")
def test_update_agent_engine_config_with_build_config(self, mock_prepare):
config = self.client.agent_engines._create_config(
mode="update",
build_config=_TEST_AGENT_ENGINE_BUILD_CONFIG,
)
assert config["spec"]["build_spec"] == {
"worker_pool": _TEST_AGENT_ENGINE_BUILD_WORKER_POOL,
"service_account": _TEST_AGENT_ENGINE_BUILD_SERVICE_ACCOUNT,
}
update_mask = config["update_mask"].split(",")
assert "spec.build_spec.worker_pool" in update_mask
assert "spec.build_spec.service_account" in update_mask

@mock.patch.object(
_agent_engines_utils,
"_create_base64_encoded_tarball",
Expand Down Expand Up @@ -2236,6 +2291,7 @@ def test_create_agent_engine_with_env_vars_dict(
agent_config_source=None,
container_spec=None,
keep_alive_probe=None,
build_config=None,
)
request_mock.assert_called_with(
"post",
Expand Down Expand Up @@ -2342,6 +2398,7 @@ def test_create_agent_engine_with_custom_service_account(
agent_config_source=None,
container_spec=None,
keep_alive_probe=None,
build_config=None,
)
request_mock.assert_called_with(
"post",
Expand Down Expand Up @@ -2447,6 +2504,7 @@ def test_create_agent_engine_with_experimental_mode(
agent_config_source=None,
container_spec=None,
keep_alive_probe=None,
build_config=None,
)
request_mock.assert_called_with(
"post",
Expand Down Expand Up @@ -2621,6 +2679,7 @@ def test_create_agent_engine_with_class_methods(
agent_config_source=None,
container_spec=None,
keep_alive_probe=None,
build_config=None,
)
request_mock.assert_called_with(
"post",
Expand Down Expand Up @@ -2721,6 +2780,7 @@ def test_create_agent_engine_with_agent_framework(
agent_config_source=None,
container_spec=None,
keep_alive_probe=None,
build_config=None,
)
request_mock.assert_called_with(
"post",
Expand Down
22 changes: 22 additions & 0 deletions vertexai/_genai/agent_engines.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2005,6 +2005,7 @@ def create(
agent_config_source=agent_config_source,
container_spec=config.container_spec,
keep_alive_probe=keep_alive_probe,
build_config=config.build_config,
)
operation = self._create(config=api_config)
reasoning_engine_id = _agent_engines_utils._get_reasoning_engine_id(
Expand Down Expand Up @@ -2318,6 +2319,7 @@ def _create_config(
container_spec: Optional[types.ReasoningEngineSpecContainerSpecDict] = None,
keep_alive_probe: Optional[dict[str, Any]] = None,
traffic_config: Optional[types.ReasoningEngineTrafficConfigDict] = None,
build_config: Optional[types.ReasoningEngineSpecBuildSpecDict] = None,
) -> types.UpdateAgentEngineConfigDict:
import sys

Expand Down Expand Up @@ -2524,6 +2526,25 @@ def _create_config(
agent_engine_spec["service_account"] = service_account
update_masks.append("spec.service_account")

if build_config is not None:
if agent_engine_spec is None:
agent_engine_spec = {}
build_spec: dict[str, Any] = {}
if isinstance(build_config, dict):
worker_pool = build_config.get("worker_pool")
build_service_account = build_config.get("service_account")
else:
worker_pool = getattr(build_config, "worker_pool", None)
build_service_account = getattr(build_config, "service_account", None)
if worker_pool is not None:
build_spec["worker_pool"] = worker_pool
update_masks.append("spec.build_spec.worker_pool")
if build_service_account is not None:
build_spec["service_account"] = build_service_account
update_masks.append("spec.build_spec.service_account")
if build_spec:
agent_engine_spec["build_spec"] = build_spec

if agent_engine_spec is not None:
config["spec"] = agent_engine_spec

Expand Down Expand Up @@ -2795,6 +2816,7 @@ def update(
container_spec=container_spec,
keep_alive_probe=keep_alive_probe,
traffic_config=traffic_config,
build_config=config.build_config,
)
operation = self._update(name=name, config=api_config)
reasoning_engine_id = _agent_engines_utils._get_reasoning_engine_id(
Expand Down
6 changes: 6 additions & 0 deletions vertexai/_genai/types/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1031,6 +1031,9 @@
from .common import ReasoningEngineRuntimeRevisionDict
from .common import ReasoningEngineRuntimeRevisionOrDict
from .common import ReasoningEngineSpec
from .common import ReasoningEngineSpecBuildSpec
from .common import ReasoningEngineSpecBuildSpecDict
from .common import ReasoningEngineSpecBuildSpecOrDict
from .common import ReasoningEngineSpecContainerSpec
from .common import ReasoningEngineSpecContainerSpecDict
from .common import ReasoningEngineSpecContainerSpecOrDict
Expand Down Expand Up @@ -2082,6 +2085,9 @@
"ReasoningEngineSpecContainerSpec",
"ReasoningEngineSpecContainerSpecDict",
"ReasoningEngineSpecContainerSpecOrDict",
"ReasoningEngineSpecBuildSpec",
"ReasoningEngineSpecBuildSpecDict",
"ReasoningEngineSpecBuildSpecOrDict",
"ReasoningEngineSpec",
"ReasoningEngineSpecDict",
"ReasoningEngineSpecOrDict",
Expand Down
Loading
Loading