Skip to content

Commit ec11bd3

Browse files
shawn-yang-googlecopybara-github
shawn-yang-google
authored and
copybara-github
committed
feat: Support for runtime controls and PSC-I in Agent Engine.
PiperOrigin-RevId: 795610914
1 parent bc7e2e9 commit ec11bd3

File tree

4 files changed

+525
-5
lines changed

4 files changed

+525
-5
lines changed

‎tests/unit/vertex_langchain/test_agent_engines.py‎

Lines changed: 306 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -589,6 +589,23 @@ def register_operations(self) -> Dict[str, List[str]]:
589589
_TEST_INSTALLATION_SCRIPT_PATH = f"{_TEST_INSTALLATION_SUBDIR}/install_package.sh"
590590

591591
_TEST_CUSTOM_SERVICE_ACCOUNT = "test-custom-service-account"
592+
_TEST_AGENT_ENGINE_PSC_INTERFACE_CONFIG = {
593+
"network_attachment": "test-network-attachment",
594+
"dns_peering_configs": [
595+
{
596+
"domain": "test-domain",
597+
"target_project": "test-target-project",
598+
"target_network": "test-target-network",
599+
}
600+
],
601+
}
602+
_TEST_AGENT_ENGINE_MIN_INSTANCES = 2
603+
_TEST_AGENT_ENGINE_MAX_INSTANCES = 4
604+
_TEST_AGENT_ENGINE_RESOURCE_LIMITS = {
605+
"cpu": "2",
606+
"memory": "4Gi",
607+
}
608+
_TEST_AGENT_ENGINE_CONTAINER_CONCURRENCY = 4
592609

593610

594611
def _create_empty_fake_package(package_name: str) -> str:
@@ -1229,6 +1246,196 @@ def test_create_agent_engine_with_service_account(
12291246
retry=_TEST_RETRY,
12301247
)
12311248

1249+
def test_create_agent_engine_with_psc_interface_config(
1250+
self,
1251+
create_agent_engine_mock,
1252+
cloud_storage_create_bucket_mock,
1253+
tarfile_open_mock,
1254+
cloudpickle_dump_mock,
1255+
cloudpickle_load_mock,
1256+
importlib_metadata_version_mock,
1257+
get_agent_engine_mock,
1258+
get_gca_resource_mock,
1259+
):
1260+
agent_engines.create(
1261+
self.test_agent,
1262+
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
1263+
requirements=_TEST_AGENT_ENGINE_REQUIREMENTS,
1264+
extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH],
1265+
psc_interface_config=_TEST_AGENT_ENGINE_PSC_INTERFACE_CONFIG,
1266+
)
1267+
test_spec = types.ReasoningEngineSpec(
1268+
package_spec=_TEST_AGENT_ENGINE_PACKAGE_SPEC,
1269+
deployment_spec=types.ReasoningEngineSpec.DeploymentSpec(
1270+
psc_interface_config=_TEST_AGENT_ENGINE_PSC_INTERFACE_CONFIG,
1271+
),
1272+
agent_framework=_agent_engines._DEFAULT_AGENT_FRAMEWORK,
1273+
)
1274+
test_spec.class_methods.append(_TEST_AGENT_ENGINE_QUERY_SCHEMA)
1275+
create_agent_engine_mock.assert_called_with(
1276+
parent=_TEST_PARENT,
1277+
reasoning_engine=types.ReasoningEngine(
1278+
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
1279+
spec=test_spec,
1280+
),
1281+
)
1282+
get_agent_engine_mock.assert_called_with(
1283+
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
1284+
retry=_TEST_RETRY,
1285+
)
1286+
1287+
def test_create_agent_engine_with_resource_limits(
1288+
self,
1289+
create_agent_engine_mock,
1290+
cloud_storage_create_bucket_mock,
1291+
tarfile_open_mock,
1292+
cloudpickle_dump_mock,
1293+
cloudpickle_load_mock,
1294+
importlib_metadata_version_mock,
1295+
get_agent_engine_mock,
1296+
get_gca_resource_mock,
1297+
):
1298+
agent_engines.create(
1299+
self.test_agent,
1300+
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
1301+
requirements=_TEST_AGENT_ENGINE_REQUIREMENTS,
1302+
extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH],
1303+
resource_limits=_TEST_AGENT_ENGINE_RESOURCE_LIMITS,
1304+
)
1305+
test_spec = types.ReasoningEngineSpec(
1306+
package_spec=_TEST_AGENT_ENGINE_PACKAGE_SPEC,
1307+
deployment_spec=types.ReasoningEngineSpec.DeploymentSpec(
1308+
resource_limits=_TEST_AGENT_ENGINE_RESOURCE_LIMITS,
1309+
),
1310+
agent_framework=_agent_engines._DEFAULT_AGENT_FRAMEWORK,
1311+
)
1312+
test_spec.class_methods.append(_TEST_AGENT_ENGINE_QUERY_SCHEMA)
1313+
create_agent_engine_mock.assert_called_with(
1314+
parent=_TEST_PARENT,
1315+
reasoning_engine=types.ReasoningEngine(
1316+
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
1317+
spec=test_spec,
1318+
),
1319+
)
1320+
get_agent_engine_mock.assert_called_with(
1321+
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
1322+
retry=_TEST_RETRY,
1323+
)
1324+
1325+
def test_create_agent_engine_with_min_instances(
1326+
self,
1327+
create_agent_engine_mock,
1328+
cloud_storage_create_bucket_mock,
1329+
tarfile_open_mock,
1330+
cloudpickle_dump_mock,
1331+
cloudpickle_load_mock,
1332+
importlib_metadata_version_mock,
1333+
get_agent_engine_mock,
1334+
get_gca_resource_mock,
1335+
):
1336+
agent_engines.create(
1337+
self.test_agent,
1338+
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
1339+
requirements=_TEST_AGENT_ENGINE_REQUIREMENTS,
1340+
extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH],
1341+
min_instances=_TEST_AGENT_ENGINE_MIN_INSTANCES,
1342+
)
1343+
test_spec = types.ReasoningEngineSpec(
1344+
package_spec=_TEST_AGENT_ENGINE_PACKAGE_SPEC,
1345+
deployment_spec=types.ReasoningEngineSpec.DeploymentSpec(
1346+
min_instances=_TEST_AGENT_ENGINE_MIN_INSTANCES,
1347+
),
1348+
agent_framework=_agent_engines._DEFAULT_AGENT_FRAMEWORK,
1349+
)
1350+
test_spec.class_methods.append(_TEST_AGENT_ENGINE_QUERY_SCHEMA)
1351+
create_agent_engine_mock.assert_called_with(
1352+
parent=_TEST_PARENT,
1353+
reasoning_engine=types.ReasoningEngine(
1354+
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
1355+
spec=test_spec,
1356+
),
1357+
)
1358+
get_agent_engine_mock.assert_called_with(
1359+
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
1360+
retry=_TEST_RETRY,
1361+
)
1362+
1363+
def test_create_agent_engine_with_max_instances(
1364+
self,
1365+
create_agent_engine_mock,
1366+
cloud_storage_create_bucket_mock,
1367+
tarfile_open_mock,
1368+
cloudpickle_dump_mock,
1369+
cloudpickle_load_mock,
1370+
importlib_metadata_version_mock,
1371+
get_agent_engine_mock,
1372+
get_gca_resource_mock,
1373+
):
1374+
agent_engines.create(
1375+
self.test_agent,
1376+
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
1377+
requirements=_TEST_AGENT_ENGINE_REQUIREMENTS,
1378+
extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH],
1379+
max_instances=_TEST_AGENT_ENGINE_MAX_INSTANCES,
1380+
)
1381+
test_spec = types.ReasoningEngineSpec(
1382+
package_spec=_TEST_AGENT_ENGINE_PACKAGE_SPEC,
1383+
deployment_spec=types.ReasoningEngineSpec.DeploymentSpec(
1384+
max_instances=_TEST_AGENT_ENGINE_MAX_INSTANCES,
1385+
),
1386+
agent_framework=_agent_engines._DEFAULT_AGENT_FRAMEWORK,
1387+
)
1388+
test_spec.class_methods.append(_TEST_AGENT_ENGINE_QUERY_SCHEMA)
1389+
create_agent_engine_mock.assert_called_with(
1390+
parent=_TEST_PARENT,
1391+
reasoning_engine=types.ReasoningEngine(
1392+
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
1393+
spec=test_spec,
1394+
),
1395+
)
1396+
get_agent_engine_mock.assert_called_with(
1397+
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
1398+
retry=_TEST_RETRY,
1399+
)
1400+
1401+
def test_create_agent_engine_with_container_concurrency(
1402+
self,
1403+
create_agent_engine_mock,
1404+
cloud_storage_create_bucket_mock,
1405+
tarfile_open_mock,
1406+
cloudpickle_dump_mock,
1407+
cloudpickle_load_mock,
1408+
importlib_metadata_version_mock,
1409+
get_agent_engine_mock,
1410+
get_gca_resource_mock,
1411+
):
1412+
agent_engines.create(
1413+
self.test_agent,
1414+
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
1415+
requirements=_TEST_AGENT_ENGINE_REQUIREMENTS,
1416+
extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH],
1417+
container_concurrency=_TEST_AGENT_ENGINE_CONTAINER_CONCURRENCY,
1418+
)
1419+
test_spec = types.ReasoningEngineSpec(
1420+
package_spec=_TEST_AGENT_ENGINE_PACKAGE_SPEC,
1421+
deployment_spec=types.ReasoningEngineSpec.DeploymentSpec(
1422+
container_concurrency=_TEST_AGENT_ENGINE_CONTAINER_CONCURRENCY,
1423+
),
1424+
agent_framework=_agent_engines._DEFAULT_AGENT_FRAMEWORK,
1425+
)
1426+
test_spec.class_methods.append(_TEST_AGENT_ENGINE_QUERY_SCHEMA)
1427+
create_agent_engine_mock.assert_called_with(
1428+
parent=_TEST_PARENT,
1429+
reasoning_engine=types.ReasoningEngine(
1430+
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
1431+
spec=test_spec,
1432+
),
1433+
)
1434+
get_agent_engine_mock.assert_called_with(
1435+
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
1436+
retry=_TEST_RETRY,
1437+
)
1438+
12321439
@pytest.mark.parametrize(
12331440
"test_case_name, test_engine_instance, expected_framework",
12341441
[
@@ -1509,6 +1716,101 @@ def test_get_agent_framework(
15091716
),
15101717
),
15111718
),
1719+
(
1720+
"Update the agent_engine with psc_interface_config attribute",
1721+
{"psc_interface_config": _TEST_AGENT_ENGINE_PSC_INTERFACE_CONFIG},
1722+
types.reasoning_engine_service.UpdateReasoningEngineRequest(
1723+
reasoning_engine=types.ReasoningEngine(
1724+
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
1725+
spec=types.ReasoningEngineSpec(
1726+
deployment_spec=types.ReasoningEngineSpec.DeploymentSpec(
1727+
psc_interface_config=_TEST_AGENT_ENGINE_PSC_INTERFACE_CONFIG,
1728+
),
1729+
),
1730+
),
1731+
update_mask=field_mask_pb2.FieldMask(
1732+
paths=[
1733+
"spec.deployment_spec.psc_interface_config",
1734+
],
1735+
),
1736+
),
1737+
),
1738+
(
1739+
"Update the agent_engine with min_instances attribute",
1740+
{"min_instances": _TEST_AGENT_ENGINE_MIN_INSTANCES},
1741+
types.reasoning_engine_service.UpdateReasoningEngineRequest(
1742+
reasoning_engine=types.ReasoningEngine(
1743+
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
1744+
spec=types.ReasoningEngineSpec(
1745+
deployment_spec=types.ReasoningEngineSpec.DeploymentSpec(
1746+
min_instances=_TEST_AGENT_ENGINE_MIN_INSTANCES,
1747+
),
1748+
),
1749+
),
1750+
update_mask=field_mask_pb2.FieldMask(
1751+
paths=[
1752+
"spec.deployment_spec.min_instances",
1753+
],
1754+
),
1755+
),
1756+
),
1757+
(
1758+
"Update the agent_engine with max_instances attribute",
1759+
{"max_instances": _TEST_AGENT_ENGINE_MAX_INSTANCES},
1760+
types.reasoning_engine_service.UpdateReasoningEngineRequest(
1761+
reasoning_engine=types.ReasoningEngine(
1762+
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
1763+
spec=types.ReasoningEngineSpec(
1764+
deployment_spec=types.ReasoningEngineSpec.DeploymentSpec(
1765+
max_instances=_TEST_AGENT_ENGINE_MAX_INSTANCES,
1766+
),
1767+
),
1768+
),
1769+
update_mask=field_mask_pb2.FieldMask(
1770+
paths=[
1771+
"spec.deployment_spec.max_instances",
1772+
],
1773+
),
1774+
),
1775+
),
1776+
(
1777+
"Update the agent_engine with resource_limits attribute",
1778+
{"resource_limits": _TEST_AGENT_ENGINE_RESOURCE_LIMITS},
1779+
types.reasoning_engine_service.UpdateReasoningEngineRequest(
1780+
reasoning_engine=types.ReasoningEngine(
1781+
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
1782+
spec=types.ReasoningEngineSpec(
1783+
deployment_spec=types.ReasoningEngineSpec.DeploymentSpec(
1784+
resource_limits=_TEST_AGENT_ENGINE_RESOURCE_LIMITS,
1785+
),
1786+
),
1787+
),
1788+
update_mask=field_mask_pb2.FieldMask(
1789+
paths=[
1790+
"spec.deployment_spec.resource_limits",
1791+
],
1792+
),
1793+
),
1794+
),
1795+
(
1796+
"Update the agent_engine with container_concurrency attribute",
1797+
{"container_concurrency": _TEST_AGENT_ENGINE_CONTAINER_CONCURRENCY},
1798+
types.reasoning_engine_service.UpdateReasoningEngineRequest(
1799+
reasoning_engine=types.ReasoningEngine(
1800+
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
1801+
spec=types.ReasoningEngineSpec(
1802+
deployment_spec=types.ReasoningEngineSpec.DeploymentSpec(
1803+
container_concurrency=_TEST_AGENT_ENGINE_CONTAINER_CONCURRENCY,
1804+
),
1805+
),
1806+
),
1807+
update_mask=field_mask_pb2.FieldMask(
1808+
paths=[
1809+
"spec.deployment_spec.container_concurrency",
1810+
],
1811+
),
1812+
),
1813+
),
15121814
],
15131815
)
15141816
def test_update_agent_engine(
@@ -2889,7 +3191,10 @@ def test_update_agent_engine_with_no_updates(
28893191
match=(
28903192
"At least one of `agent_engine`, `requirements`, "
28913193
"`extra_packages`, `display_name`, `description`, "
2892-
"`env_vars`, or `build_options` must be specified."
3194+
"`env_vars`, `build_options`, `service_account`, "
3195+
"`psc_interface_config`, `min_instances`, `max_instances`, "
3196+
"`resource_limits`, or `container_concurrency` "
3197+
"must be specified."
28933198
),
28943199
):
28953200
test_agent_engine = _generate_agent_engine_to_update()

0 commit comments

Comments
 (0)