From 8ef6a4bde05b4ae2ee6671157e408ab8ff99cc85 Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Wed, 25 Mar 2026 08:18:41 -0700 Subject: [PATCH] feat: add support for keep alive probe in agent engines Keep alive probe allows reasoning engine users to configure a probe that a deployment host can use to keep the container alive, based on the probe settings. If the keep alive endpoint returns a 2xx status, the deployment host will make a best effort (up to 1 hour) to keep the container alive. Reasoning engine users with custom container specs (BYOC) have the option to configure a custom keep alive probe while the users without custom container specs (BYOC) have the option to configure an empty keep alive probe {} and the reasoning engine platform will handle the configuration and logic for keep alive probe. To opt in, users should set the keep alive probe field when creating or updating reasoning engines. PiperOrigin-RevId: 889254469 --- .../unit/vertexai/genai/test_agent_engines.py | 165 ++++++++++++++++++ vertexai/_genai/agent_engines.py | 27 ++- vertexai/_genai/types/__init__.py | 12 ++ vertexai/_genai/types/common.py | 92 ++++++++++ 4 files changed, 292 insertions(+), 4 deletions(-) diff --git a/tests/unit/vertexai/genai/test_agent_engines.py b/tests/unit/vertexai/genai/test_agent_engines.py index 99184656ae..a8bf3a6d60 100644 --- a/tests/unit/vertexai/genai/test_agent_engines.py +++ b/tests/unit/vertexai/genai/test_agent_engines.py @@ -538,6 +538,12 @@ 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_KEEP_ALIVE_PROBE = { + "http_get": { + "path": "/health", + }, + "max_seconds": 60, +} _TEST_AGENT_ENGINE_SPEC = _genai_types.ReasoningEngineSpecDict( agent_framework=_TEST_AGENT_ENGINE_FRAMEWORK, class_methods=[_TEST_AGENT_ENGINE_CLASS_METHOD_1], @@ -1071,6 +1077,7 @@ def test_create_agent_engine_config_with_source_packages( config["spec"]["identity_type"] == _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT ) + assert "keep_alive_probe" not in config["spec"].get("deployment_spec", {}) def test_create_agent_engine_config_with_developer_connect_source(self): with tempfile.TemporaryDirectory() as tmpdir: @@ -1112,6 +1119,29 @@ def test_create_agent_engine_config_with_developer_connect_source(self): config["spec"]["identity_type"] == _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT ) + assert "keep_alive_probe" not in config["spec"].get("deployment_spec", {}) + + @mock.patch.object( + _agent_engines_utils, + "_create_base64_encoded_tarball", + return_value="test_tarball", + ) + def test_create_agent_engine_config_with_empty_keep_alive_probe( + self, mock_create_base64_encoded_tarball + ): + with tempfile.TemporaryDirectory() as tmpdir: + test_file_path = os.path.join(tmpdir, "test_file.txt") + with open(test_file_path, "w") as f: + f.write("test content") + config = self.client.agent_engines._create_config( + mode="create", + source_packages=[test_file_path], + class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS, + entrypoint_module="main", + entrypoint_object="app", + keep_alive_probe={}, + ) + assert "keep_alive_probe" in config["spec"].get("deployment_spec", {}) def test_create_agent_engine_config_with_agent_config_source_and_requirements_file( self, @@ -1321,6 +1351,33 @@ def test_create_agent_engine_config_with_container_spec(self): config["spec"]["identity_type"] == _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT ) + assert "keep_alive_probe" not in config["spec"].get("deployment_spec", {}) + + def test_create_agent_engine_config_with_container_spec_and_keep_alive_probe( + self, + ): + container_spec = {"image_uri": "gcr.io/test-project/test-image"} + config = self.client.agent_engines._create_config( + mode="create", + display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, + description=_TEST_AGENT_ENGINE_DESCRIPTION, + container_spec=container_spec, + class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS, + identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT, + keep_alive_probe=_TEST_AGENT_ENGINE_KEEP_ALIVE_PROBE, + ) + assert config["display_name"] == _TEST_AGENT_ENGINE_DISPLAY_NAME + assert config["description"] == _TEST_AGENT_ENGINE_DESCRIPTION + assert config["spec"]["container_spec"] == container_spec + assert config["spec"]["class_methods"] == _TEST_AGENT_ENGINE_CLASS_METHODS + assert ( + config["spec"]["identity_type"] + == _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT + ) + assert ( + config["spec"]["deployment_spec"]["keep_alive_probe"] + == _TEST_AGENT_ENGINE_KEEP_ALIVE_PROBE + ) def test_create_agent_engine_config_with_container_spec_and_others_raises(self): container_spec = {"image_uri": "gcr.io/test-project/test-image"} @@ -2116,6 +2173,7 @@ def test_create_agent_engine_with_env_vars_dict( image_spec=None, agent_config_source=None, container_spec=None, + keep_alive_probe=None, ) request_mock.assert_called_with( "post", @@ -2220,6 +2278,7 @@ def test_create_agent_engine_with_custom_service_account( image_spec=None, agent_config_source=None, container_spec=None, + keep_alive_probe=None, ) request_mock.assert_called_with( "post", @@ -2323,6 +2382,7 @@ def test_create_agent_engine_with_experimental_mode( image_spec=None, agent_config_source=None, container_spec=None, + keep_alive_probe=None, ) request_mock.assert_called_with( "post", @@ -2495,6 +2555,7 @@ def test_create_agent_engine_with_class_methods( image_spec=None, agent_config_source=None, container_spec=None, + keep_alive_probe=None, ) request_mock.assert_called_with( "post", @@ -2593,6 +2654,7 @@ def test_create_agent_engine_with_agent_framework( image_spec=None, agent_config_source=None, container_spec=None, + keep_alive_probe=None, ) request_mock.assert_called_with( "post", @@ -2795,6 +2857,109 @@ def test_update_agent_engine_env_vars( None, ) + @mock.patch.object(_agent_engines_utils, "_prepare") + @mock.patch.object(_agent_engines_utils, "_await_operation") + def test_update_agent_engine_with_empty_keep_alive_probe( + self, mock_await_operation, mock_prepare + ): + 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.update( + name=_TEST_AGENT_ENGINE_RESOURCE_NAME, + agent=self.test_agent, + config=_genai_types.AgentEngineConfig( + staging_bucket=_TEST_STAGING_BUCKET, + keep_alive_probe={}, + ), + ) + update_mask = ",".join( + [ + "spec.package_spec.pickle_object_gcs_uri", + "spec.package_spec.requirements_gcs_uri", + "spec.class_methods", + "spec.deployment_spec.keep_alive_probe", + "spec.agent_framework", + ] + ) + query_params = {"updateMask": update_mask} + request_mock.assert_called_with( + "patch", + f"{_TEST_AGENT_ENGINE_RESOURCE_NAME}?{urlencode(query_params)}", + { + "_url": {"name": _TEST_AGENT_ENGINE_RESOURCE_NAME}, + "spec": { + "agent_framework": _TEST_AGENT_ENGINE_FRAMEWORK, + "class_methods": mock.ANY, + "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, + }, + "deployment_spec": {"keep_alive_probe": {}}, + }, + "_query": {"updateMask": update_mask}, + }, + None, + ) + + @mock.patch.object(_agent_engines_utils, "_await_operation") + def test_update_agent_engine_with_container_spec_and_keep_alive_probe( + self, mock_await_operation + ): + mock_await_operation.return_value = _genai_types.AgentEngineOperation( + response=_genai_types.ReasoningEngine( + name=_TEST_AGENT_ENGINE_RESOURCE_NAME, + spec=_TEST_AGENT_ENGINE_SPEC, + ) + ) + container_spec = {"image_uri": "gcr.io/test-project/test-image"} + 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.update( + name=_TEST_AGENT_ENGINE_RESOURCE_NAME, + config=_genai_types.AgentEngineConfig( + container_spec=container_spec, + keep_alive_probe=_TEST_AGENT_ENGINE_KEEP_ALIVE_PROBE, + class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS, + ), + ) + update_mask = ",".join( + [ + "spec.class_methods", + "spec.container_spec", + "spec.deployment_spec.keep_alive_probe", + "spec.agent_framework", + ] + ) + query_params = {"updateMask": update_mask} + request_mock.assert_called_with( + "patch", + f"{_TEST_AGENT_ENGINE_RESOURCE_NAME}?{urlencode(query_params)}", + { + "_url": {"name": _TEST_AGENT_ENGINE_RESOURCE_NAME}, + "spec": { + "agent_framework": "custom", + "container_spec": container_spec, + "deployment_spec": { + "keep_alive_probe": _TEST_AGENT_ENGINE_KEEP_ALIVE_PROBE, + }, + "class_methods": mock.ANY, + }, + "_query": {"updateMask": update_mask}, + }, + None, + ) + @mock.patch.object(_agent_engines_utils, "_await_operation") def test_update_agent_engine_display_name(self, mock_await_operation): mock_await_operation.return_value = _genai_types.AgentEngineOperation( diff --git a/vertexai/_genai/agent_engines.py b/vertexai/_genai/agent_engines.py index d41a6f7280..c2472ff381 100644 --- a/vertexai/_genai/agent_engines.py +++ b/vertexai/_genai/agent_engines.py @@ -1317,6 +1317,11 @@ def create( agent_config_source = config.agent_config_source if agent_config_source is not None: agent_config_source = json.loads(agent_config_source.model_dump_json()) + keep_alive_probe = config.keep_alive_probe + if keep_alive_probe is not None: + keep_alive_probe = json.loads( + keep_alive_probe.model_dump_json(exclude_none=True) + ) if agent and agent_engine: raise ValueError("Please specify only one of `agent` or `agent_engine`.") elif agent_engine: @@ -1357,6 +1362,7 @@ def create( image_spec=config.image_spec, agent_config_source=agent_config_source, container_spec=config.container_spec, + keep_alive_probe=keep_alive_probe, ) operation = self._create(config=api_config) reasoning_engine_id = _agent_engines_utils._get_reasoning_engine_id( @@ -1665,6 +1671,7 @@ def _create_config( types.ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict ] = None, container_spec: Optional[types.ReasoningEngineSpecContainerSpecDict] = None, + keep_alive_probe: Optional[dict[str, Any]] = None, ) -> types.UpdateAgentEngineConfigDict: import sys @@ -1794,14 +1801,15 @@ def _create_config( or max_instances is not None or resource_limits is not None or container_concurrency is not None + or keep_alive_probe 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`, or `container_concurrency`, " - "you must also provide the `agent` variable or the source code " - "options (`source_packages`, `developer_connect_source` or " - "`agent_config_source`)." + "`max_instances`, `resource_limits`, `container_concurrency`, or " + "`keep_alive_probe`, you must also provide the `agent` variable or " + "the source code options (`source_packages`, " + "`developer_connect_source` or `agent_config_source`)." ) if agent_engine_spec is not None: @@ -1816,6 +1824,7 @@ def _create_config( max_instances=max_instances, resource_limits=resource_limits, container_concurrency=container_concurrency, + keep_alive_probe=keep_alive_probe, ) update_masks.extend(deployment_update_masks) agent_engine_spec["deployment_spec"] = deployment_spec @@ -1878,6 +1887,7 @@ def _generate_deployment_spec_or_raise( max_instances: Optional[int] = None, resource_limits: Optional[dict[str, str]] = None, container_concurrency: Optional[int] = None, + keep_alive_probe: Optional[dict[str, Any]] = None, ) -> Tuple[dict[str, Any], Sequence[str]]: deployment_spec: dict[str, Any] = {} update_masks = [] @@ -1925,6 +1935,9 @@ def _generate_deployment_spec_or_raise( if container_concurrency: deployment_spec["container_concurrency"] = container_concurrency update_masks.append("spec.deployment_spec.container_concurrency") + if keep_alive_probe is not None: + deployment_spec["keep_alive_probe"] = keep_alive_probe + update_masks.append("spec.deployment_spec.keep_alive_probe") return deployment_spec, update_masks def _update_deployment_spec_with_env_vars_dict_or_raise( @@ -2066,6 +2079,11 @@ def update( agent_config_source = config.agent_config_source if agent_config_source is not None: agent_config_source = json.loads(agent_config_source.model_dump_json()) + keep_alive_probe = config.keep_alive_probe + if keep_alive_probe is not None: + keep_alive_probe = json.loads( + keep_alive_probe.model_dump_json(exclude_none=True) + ) if agent and agent_engine: raise ValueError("Please specify only one of `agent` or `agent_engine`.") elif agent_engine: @@ -2112,6 +2130,7 @@ def update( image_spec=image_spec, agent_config_source=agent_config_source, container_spec=container_spec, + keep_alive_probe=keep_alive_probe, ) operation = self._update(name=name, config=api_config) reasoning_engine_id = _agent_engines_utils._get_reasoning_engine_id( diff --git a/vertexai/_genai/types/__init__.py b/vertexai/_genai/types/__init__.py index ef47dd9d35..8f56481e3e 100644 --- a/vertexai/_genai/types/__init__.py +++ b/vertexai/_genai/types/__init__.py @@ -522,6 +522,12 @@ from .common import IntermediateExtractedMemoryDict from .common import IntermediateExtractedMemoryOrDict from .common import JobState +from .common import KeepAliveProbe +from .common import KeepAliveProbeDict +from .common import KeepAliveProbeHttpGet +from .common import KeepAliveProbeHttpGetDict +from .common import KeepAliveProbeHttpGetOrDict +from .common import KeepAliveProbeOrDict from .common import Language from .common import ListAgentEngineConfig from .common import ListAgentEngineConfigDict @@ -1603,6 +1609,12 @@ "SecretEnvVar", "SecretEnvVarDict", "SecretEnvVarOrDict", + "KeepAliveProbeHttpGet", + "KeepAliveProbeHttpGetDict", + "KeepAliveProbeHttpGetOrDict", + "KeepAliveProbe", + "KeepAliveProbeDict", + "KeepAliveProbeOrDict", "ReasoningEngineSpecDeploymentSpec", "ReasoningEngineSpecDeploymentSpecDict", "ReasoningEngineSpecDeploymentSpecOrDict", diff --git a/vertexai/_genai/types/common.py b/vertexai/_genai/types/common.py index e6a7ea4b3c..79ec22b252 100644 --- a/vertexai/_genai/types/common.py +++ b/vertexai/_genai/types/common.py @@ -6501,6 +6501,58 @@ class SecretEnvVarDict(TypedDict, total=False): SecretEnvVarOrDict = Union[SecretEnvVar, SecretEnvVarDict] +class KeepAliveProbeHttpGet(_common.BaseModel): + """Specifies the HTTP GET configuration for the probe.""" + + path: Optional[str] = Field( + default=None, + description="""Required. Specifies the path of the HTTP GET request (e.g., "/is_busy").""", + ) + port: Optional[int] = Field( + default=None, + description="""Optional. Specifies the port number on the container to which the request is sent.""", + ) + + +class KeepAliveProbeHttpGetDict(TypedDict, total=False): + """Specifies the HTTP GET configuration for the probe.""" + + path: Optional[str] + """Required. Specifies the path of the HTTP GET request (e.g., "/is_busy").""" + + port: Optional[int] + """Optional. Specifies the port number on the container to which the request is sent.""" + + +KeepAliveProbeHttpGetOrDict = Union[KeepAliveProbeHttpGet, KeepAliveProbeHttpGetDict] + + +class KeepAliveProbe(_common.BaseModel): + """Represents 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.""" + + http_get: Optional[KeepAliveProbeHttpGet] = Field( + default=None, + description="""Optional. Specifies the HTTP GET configuration for the probe.""", + ) + max_seconds: Optional[int] = Field( + default=None, + description="""Optional. Specifies the maximum duration (in seconds) to keep the instance alive via this probe. Can be a maximum of 3600 seconds (1 hour).""", + ) + + +class KeepAliveProbeDict(TypedDict, total=False): + """Represents 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.""" + + http_get: Optional[KeepAliveProbeHttpGetDict] + """Optional. Specifies the HTTP GET configuration for the probe.""" + + max_seconds: Optional[int] + """Optional. Specifies the maximum duration (in seconds) to keep the instance alive via this probe. Can be a maximum of 3600 seconds (1 hour).""" + + +KeepAliveProbeOrDict = Union[KeepAliveProbe, KeepAliveProbeDict] + + class ReasoningEngineSpecDeploymentSpec(_common.BaseModel): """The specification of a Reasoning Engine deployment.""" @@ -6534,6 +6586,10 @@ class ReasoningEngineSpecDeploymentSpec(_common.BaseModel): default=None, description="""Optional. Environment variables where the value is a secret in Cloud Secret Manager. To use this feature, add 'Secret Manager Secret Accessor' role (roles/secretmanager.secretAccessor) to AI Platform Reasoning Engine Service Agent.""", ) + keep_alive_probe: Optional[KeepAliveProbe] = Field( + 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.""", + ) class ReasoningEngineSpecDeploymentSpecDict(TypedDict, total=False): @@ -6563,6 +6619,9 @@ class ReasoningEngineSpecDeploymentSpecDict(TypedDict, total=False): secret_env: Optional[list[SecretEnvVarDict]] """Optional. Environment variables where the value is a secret in Cloud Secret Manager. To use this feature, add 'Secret Manager Secret Accessor' role (roles/secretmanager.secretAccessor) to AI Platform Reasoning Engine Service Agent.""" + 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.""" + ReasoningEngineSpecDeploymentSpecOrDict = Union[ ReasoningEngineSpecDeploymentSpec, ReasoningEngineSpecDeploymentSpecDict @@ -7251,6 +7310,12 @@ class CreateAgentEngineConfig(_common.BaseModel): subdirectory and the path must be added to `extra_packages`. """, ) + keep_alive_probe: Optional[KeepAliveProbe] = Field( + 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.""", + ) class CreateAgentEngineConfigDict(TypedDict, total=False): @@ -7382,6 +7447,11 @@ class CreateAgentEngineConfigDict(TypedDict, total=False): subdirectory and the path must be added to `extra_packages`. """ + 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.""" + CreateAgentEngineConfigOrDict = Union[ CreateAgentEngineConfig, CreateAgentEngineConfigDict @@ -7898,6 +7968,12 @@ class UpdateAgentEngineConfig(_common.BaseModel): subdirectory and the path must be added to `extra_packages`. """, ) + keep_alive_probe: Optional[KeepAliveProbe] = Field( + 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.""", + ) update_mask: Optional[str] = Field( default=None, description="""The update mask to apply. For the `FieldMask` definition, see @@ -8034,6 +8110,11 @@ class UpdateAgentEngineConfigDict(TypedDict, total=False): subdirectory and the path must be added to `extra_packages`. """ + 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.""" + update_mask: Optional[str] """The update mask to apply. For the `FieldMask` definition, see https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask.""" @@ -15515,6 +15596,12 @@ class AgentEngineConfig(_common.BaseModel): container_spec: Optional[ReasoningEngineSpecContainerSpec] = Field( default=None, description="""The container spec for the Agent Engine.""" ) + keep_alive_probe: Optional[KeepAliveProbe] = Field( + 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.""", + ) class AgentEngineConfigDict(TypedDict, total=False): @@ -15689,6 +15776,11 @@ class AgentEngineConfigDict(TypedDict, total=False): container_spec: Optional[ReasoningEngineSpecContainerSpecDict] """The container spec for the Agent Engine.""" + 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.""" + AgentEngineConfigOrDict = Union[AgentEngineConfig, AgentEngineConfigDict]