Skip to content

feat: add dedicated ingress endpoint configuration to reasoning engine #6762

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
+163 −2
Open

feat: add dedicated ingress endpoint configuration to reasoning engine #6762

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
120 changes: 120 additions & 0 deletions tests/unit/vertexai/genai/test_agent_engines.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1002,6 +1002,7 @@ def test_create_agent_engine_config_full(self, mock_prepare):
container_concurrency=_TEST_AGENT_ENGINE_CONTAINER_CONCURRENCY,
encryption_spec=_TEST_AGENT_ENGINE_ENCRYPTION_SPEC,
python_version=_TEST_PYTHON_VERSION_OVERRIDE,
dedicated_ingress_endpoint_enabled=True,
)
assert config["display_name"] == _TEST_AGENT_ENGINE_DISPLAY_NAME
assert config["description"] == _TEST_AGENT_ENGINE_DESCRIPTION
Expand Down Expand Up @@ -1032,6 +1033,7 @@ def test_create_agent_engine_config_full(self, mock_prepare):
"max_instances": _TEST_AGENT_ENGINE_MAX_INSTANCES,
"resource_limits": _TEST_AGENT_ENGINE_RESOURCE_LIMITS,
"container_concurrency": _TEST_AGENT_ENGINE_CONTAINER_CONCURRENCY,
"dedicated_ingress_endpoint_enabled": True,
}
assert config["encryption_spec"] == _TEST_AGENT_ENGINE_ENCRYPTION_SPEC
assert config["spec"]["class_methods"] == [_TEST_AGENT_ENGINE_CLASS_METHOD_1]
Expand Down Expand Up @@ -1486,6 +1488,7 @@ def test_update_agent_engine_config_full(self, mock_prepare):
service_account=_TEST_AGENT_ENGINE_CUSTOM_SERVICE_ACCOUNT,
identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT,
python_version=_TEST_PYTHON_VERSION_OVERRIDE,
dedicated_ingress_endpoint_enabled=True,
)
assert config["display_name"] == _TEST_AGENT_ENGINE_DISPLAY_NAME
assert config["description"] == _TEST_AGENT_ENGINE_DESCRIPTION
Expand All @@ -1510,6 +1513,7 @@ def test_update_agent_engine_config_full(self, mock_prepare):
},
},
],
"dedicated_ingress_endpoint_enabled": True,
}
assert config["spec"]["class_methods"] == [_TEST_AGENT_ENGINE_CLASS_METHOD_1]
assert (
Expand All @@ -1530,6 +1534,7 @@ def test_update_agent_engine_config_full(self, mock_prepare):
"spec.class_methods",
"spec.deployment_spec.env",
"spec.deployment_spec.secret_env",
"spec.deployment_spec.dedicated_ingress_endpoint_enabled",
"spec.agent_framework",
"spec.identity_type",
"spec.service_account",
Expand Down Expand Up @@ -2190,6 +2195,7 @@ def test_create_agent_engine_with_env_vars_dict(
image_spec=None,
agent_config_source=None,
container_spec=None,
dedicated_ingress_endpoint_enabled=None,
keep_alive_probe=None,
)
request_mock.assert_called_with(
Expand All @@ -2213,6 +2219,116 @@ def test_create_agent_engine_with_env_vars_dict(
None,
)

@mock.patch.object(agent_engines.AgentEngines, "_create_config")
@mock.patch.object(_agent_engines_utils, "_await_operation")
@mock.patch.object(
_agent_engines_utils,
"_get_reasoning_engine_id",
return_value=_TEST_RESOURCE_ID,
)
def test_create_agent_engine_with_dedicated_ingress_endpoint(
self,
mock_get_reasoning_engine_id,
mock_await_operation,
mock_create_config,
):
mock_create_config.return_value = {
"display_name": _TEST_AGENT_ENGINE_DISPLAY_NAME,
"description": _TEST_AGENT_ENGINE_DESCRIPTION,
"spec": {
"package_spec": {
"python_version": _TEST_PYTHON_VERSION,
"pickle_object_gcs_uri": _TEST_AGENT_ENGINE_GCS_URI,
"requirements_gcs_uri": _TEST_AGENT_ENGINE_REQUIREMENTS_GCS_URI,
},
"class_methods": [_TEST_AGENT_ENGINE_CLASS_METHOD_1],
"deployment_spec": {
"dedicated_ingress_endpoint_enabled": True,
},
"agent_framework": _TEST_AGENT_ENGINE_FRAMEWORK,
},
}
mock_await_operation.return_value = _genai_types.AgentEngineOperation(
response=_genai_types.ReasoningEngine(
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
spec=_TEST_AGENT_ENGINE_SPEC,
)
)
with mock.patch.object(
self.client.agent_engines._api_client, "request"
) as request_mock:
request_mock.return_value = genai_types.HttpResponse(body="")
self.client.agent_engines.create(
agent=self.test_agent,
config=_genai_types.AgentEngineConfig(
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
requirements=_TEST_AGENT_ENGINE_REQUIREMENTS,
extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH],
staging_bucket=_TEST_STAGING_BUCKET,
dedicated_ingress_endpoint_enabled=True,
),
)
mock_create_config.assert_called_with(
mode="create",
agent=self.test_agent,
staging_bucket=_TEST_STAGING_BUCKET,
requirements=_TEST_AGENT_ENGINE_REQUIREMENTS,
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
description=None,
gcs_dir_name=None,
extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH],
env_vars=None,
service_account=None,
identity_type=None,
context_spec=None,
psc_interface_config=None,
agent_gateway_config=None,
min_instances=None,
max_instances=None,
resource_limits=None,
container_concurrency=None,
encryption_spec=None,
agent_server_mode=None,
labels=None,
class_methods=None,
source_packages=None,
developer_connect_source=None,
entrypoint_module=None,
entrypoint_object=None,
requirements_file=None,
agent_framework=None,
python_version=None,
build_options=None,
image_spec=None,
agent_config_source=None,
container_spec=None,
dedicated_ingress_endpoint_enabled=True,
keep_alive_probe=None,
)
request_mock.assert_called_with(
"post",
"reasoningEngines",
{
"displayName": _TEST_AGENT_ENGINE_DISPLAY_NAME,
"description": _TEST_AGENT_ENGINE_DESCRIPTION,
"spec": {
"agent_framework": _TEST_AGENT_ENGINE_FRAMEWORK,
"class_methods": [_TEST_AGENT_ENGINE_CLASS_METHOD_1],
"deployment_spec": {
"dedicated_ingress_endpoint_enabled": True,
},
"package_spec": {
"pickle_object_gcs_uri": _TEST_AGENT_ENGINE_GCS_URI,
"python_version": _TEST_PYTHON_VERSION,
"requirements_gcs_uri": (
_TEST_AGENT_ENGINE_REQUIREMENTS_GCS_URI
),
},
},
},
None,
)

@mock.patch.object(agent_engines.AgentEngines, "_create_config")
@mock.patch.object(_agent_engines_utils, "_await_operation")
@mock.patch.object(
Expand Down Expand Up @@ -2296,6 +2412,7 @@ def test_create_agent_engine_with_custom_service_account(
image_spec=None,
agent_config_source=None,
container_spec=None,
dedicated_ingress_endpoint_enabled=None,
keep_alive_probe=None,
)
request_mock.assert_called_with(
Expand Down Expand Up @@ -2401,6 +2518,7 @@ def test_create_agent_engine_with_experimental_mode(
image_spec=None,
agent_config_source=None,
container_spec=None,
dedicated_ingress_endpoint_enabled=None,
keep_alive_probe=None,
)
request_mock.assert_called_with(
Expand Down Expand Up @@ -2575,6 +2693,7 @@ def test_create_agent_engine_with_class_methods(
image_spec=None,
agent_config_source=None,
container_spec=None,
dedicated_ingress_endpoint_enabled=None,
keep_alive_probe=None,
)
request_mock.assert_called_with(
Expand Down Expand Up @@ -2675,6 +2794,7 @@ def test_create_agent_engine_with_agent_framework(
image_spec=None,
agent_config_source=None,
container_spec=None,
dedicated_ingress_endpoint_enabled=None,
keep_alive_probe=None,
)
request_mock.assert_called_with(
Expand Down
17 changes: 15 additions & 2 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,
dedicated_ingress_endpoint_enabled=config.dedicated_ingress_endpoint_enabled,
)
operation = self._create(config=api_config)
reasoning_engine_id = _agent_engines_utils._get_reasoning_engine_id(
Expand Down Expand Up @@ -2316,6 +2317,7 @@ def _create_config(
types.ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict
] = None,
container_spec: Optional[types.ReasoningEngineSpecContainerSpecDict] = None,
dedicated_ingress_endpoint_enabled: Optional[bool] = None,
keep_alive_probe: Optional[dict[str, Any]] = None,
traffic_config: Optional[types.ReasoningEngineTrafficConfigDict] = None,
) -> types.UpdateAgentEngineConfigDict:
Expand Down Expand Up @@ -2452,12 +2454,13 @@ def _create_config(
or resource_limits is not None
or container_concurrency is not None
or keep_alive_probe is not None
or dedicated_ingress_endpoint_enabled is not None
)
if agent_engine_spec is None and is_deployment_spec_updated:
raise ValueError(
"To update `env_vars`, `psc_interface_config`, `min_instances`, "
"`max_instances`, `resource_limits`, `container_concurrency`, or "
"`keep_alive_probe`, you must also provide the `agent` variable or "
"`max_instances`, `resource_limits`, `container_concurrency`, "
"`keep_alive_probe`, or `dedicated_ingress_endpoint_enabled`, you must also provide the `agent` variable or "
"the source code options (`source_packages`, "
"`developer_connect_source` or `agent_config_source`)."
)
Expand All @@ -2476,6 +2479,7 @@ def _create_config(
resource_limits=resource_limits,
container_concurrency=container_concurrency,
keep_alive_probe=keep_alive_probe,
dedicated_ingress_endpoint_enabled=dedicated_ingress_endpoint_enabled,
)
update_masks.extend(deployment_update_masks)
agent_engine_spec["deployment_spec"] = deployment_spec
Expand Down Expand Up @@ -2544,6 +2548,7 @@ def _generate_deployment_spec_or_raise(
resource_limits: Optional[dict[str, str]] = None,
container_concurrency: Optional[int] = None,
keep_alive_probe: Optional[dict[str, Any]] = None,
dedicated_ingress_endpoint_enabled: Optional[bool] = None,
) -> Tuple[dict[str, Any], Sequence[str]]:
deployment_spec: dict[str, Any] = {}
update_masks = []
Expand Down Expand Up @@ -2597,6 +2602,13 @@ def _generate_deployment_spec_or_raise(
if keep_alive_probe is not None:
deployment_spec["keep_alive_probe"] = keep_alive_probe
update_masks.append("spec.deployment_spec.keep_alive_probe")
if dedicated_ingress_endpoint_enabled is not None:
deployment_spec["dedicated_ingress_endpoint_enabled"] = (
dedicated_ingress_endpoint_enabled
)
update_masks.append(
"spec.deployment_spec.dedicated_ingress_endpoint_enabled"
)
return deployment_spec, update_masks

def _update_deployment_spec_with_env_vars_dict_or_raise(
Expand Down Expand Up @@ -2794,6 +2806,7 @@ def update(
agent_config_source=agent_config_source,
container_spec=container_spec,
keep_alive_probe=keep_alive_probe,
dedicated_ingress_endpoint_enabled=config.dedicated_ingress_endpoint_enabled,
traffic_config=traffic_config,
)
operation = self._update(name=name, config=api_config)
Expand Down
28 changes: 28 additions & 0 deletions vertexai/_genai/types/common.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7629,6 +7629,10 @@ class ReasoningEngineSpecDeploymentSpec(_common.BaseModel):
default=None,
description="""Optional. Specifies the configuration for keep-alive probe. Contains configuration on a specified endpoint that a deployment host should use to keep the container alive based on the probe settings.""",
)
dedicated_ingress_endpoint_enabled: Optional[bool] = Field(
default=None,
description="""Optional. If true, the Reasoning Engine will be deployed with a dedicated ingress endpoint.""",
)


class ReasoningEngineSpecDeploymentSpecDict(TypedDict, total=False):
Expand Down Expand Up @@ -7666,6 +7670,9 @@ class ReasoningEngineSpecDeploymentSpecDict(TypedDict, total=False):
keep_alive_probe: Optional[KeepAliveProbeDict]
"""Optional. Specifies the configuration for keep-alive probe. Contains configuration on a specified endpoint that a deployment host should use to keep the container alive based on the probe settings."""

dedicated_ingress_endpoint_enabled: Optional[bool]
"""Optional. If true, the Reasoning Engine will be deployed with a dedicated ingress endpoint."""


ReasoningEngineSpecDeploymentSpecOrDict = Union[
ReasoningEngineSpecDeploymentSpec, ReasoningEngineSpecDeploymentSpecDict
Expand Down Expand Up @@ -8476,6 +8483,10 @@ class CreateAgentEngineConfig(_common.BaseModel):
Contains configuration on a specified endpoint that a deployment host
should use to keep the container alive based on the probe settings.""",
)
dedicated_ingress_endpoint_enabled: Optional[bool] = Field(
default=None,
description="""Optional. If true, the Reasoning Engine will be deployed with a dedicated ingress endpoint.""",
)


class CreateAgentEngineConfigDict(TypedDict, total=False):
Expand Down Expand Up @@ -8617,6 +8628,9 @@ class CreateAgentEngineConfigDict(TypedDict, total=False):
Contains configuration on a specified endpoint that a deployment host
should use to keep the container alive based on the probe settings."""

dedicated_ingress_endpoint_enabled: Optional[bool]
"""Optional. If true, the Reasoning Engine will be deployed with a dedicated ingress endpoint."""


CreateAgentEngineConfigOrDict = Union[
CreateAgentEngineConfig, CreateAgentEngineConfigDict
Expand Down Expand Up @@ -9143,6 +9157,10 @@ class UpdateAgentEngineConfig(_common.BaseModel):
Contains configuration on a specified endpoint that a deployment host
should use to keep the container alive based on the probe settings.""",
)
dedicated_ingress_endpoint_enabled: Optional[bool] = Field(
default=None,
description="""Optional. If true, the Reasoning Engine will be deployed with a dedicated ingress endpoint.""",
)
update_mask: Optional[str] = Field(
default=None,
description="""The update mask to apply. For the `FieldMask` definition, see
Expand Down Expand Up @@ -9293,6 +9311,9 @@ class UpdateAgentEngineConfigDict(TypedDict, total=False):
Contains configuration on a specified endpoint that a deployment host
should use to keep the container alive based on the probe settings."""

dedicated_ingress_endpoint_enabled: Optional[bool]
"""Optional. If true, the Reasoning Engine will be deployed with a dedicated ingress endpoint."""

update_mask: Optional[str]
"""The update mask to apply. For the `FieldMask` definition, see
https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask."""
Expand Down Expand Up @@ -19574,6 +19595,10 @@ class AgentEngineConfig(_common.BaseModel):
traffic_config: Optional[ReasoningEngineTrafficConfig] = Field(
default=None, description="""The traffic config for the Agent Engine."""
)
dedicated_ingress_endpoint_enabled: Optional[bool] = Field(
default=None,
description="""Optional. If true, the Reasoning Engine will be deployed with a dedicated ingress endpoint.""",
)


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

dedicated_ingress_endpoint_enabled: Optional[bool]
"""Optional. If true, the Reasoning Engine will be deployed with a dedicated ingress endpoint."""


AgentEngineConfigOrDict = Union[AgentEngineConfig, AgentEngineConfigDict]

Expand Down
Loading