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

Author: vertex-sdk-bot

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" not 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(

vertexai/_genai/agent_engines.py

@@ -1317,6 +1317,9 @@ 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())
         if agent and agent_engine:
             raise ValueError("Please specify only one of `agent` or `agent_engine`.")
         elif agent_engine:
@@ -1357,6 +1360,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(
@@ -1663,6 +1667,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
 
@@ -1795,11 +1800,12 @@ 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`, "
+                "`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`)."
@@ -1817,6 +1823,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
@@ -1879,6 +1886,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 = []
@@ -1926,6 +1934,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:
+            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(
@@ -2067,6 +2078,9 @@ 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())
         if agent and agent_engine:
             raise ValueError("Please specify only one of `agent` or `agent_engine`.")
         elif agent_engine:
@@ -2113,6 +2127,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(

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
@@ -1645,6 +1651,12 @@
     "AgentEngineOperation",
     "AgentEngineOperationDict",
     "AgentEngineOperationOrDict",
+    "KeepAliveProbeHttpGet",
+    "KeepAliveProbeHttpGetDict",
+    "KeepAliveProbeHttpGetOrDict",
+    "KeepAliveProbe",
+    "KeepAliveProbeDict",
+    "KeepAliveProbeOrDict",
     "CreateAgentEngineConfig",
     "CreateAgentEngineConfigDict",
     "CreateAgentEngineConfigOrDict",