Skip to content

fix: allow clearing env_vars / secret_env and updating deployment_spec without agent or source options in runtimes.update() and agent_engines.update() #7300

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

fix: allow clearing env_vars / secret_env and updating deployment_spec without agent or source options in runtimes.update() and agent_engines.update() #7300

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
54 changes: 33 additions & 21 deletions agentplatform/_genai/runtimes.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2277,6 +2277,7 @@ def _create_config(
update_masks.append("traffic_config")
config["traffic_config"] = traffic_config

caller_provided_env_vars = env_vars is not None
if agent_framework == "google-adk":
env_vars = _runtimes_utils._add_telemetry_enablement_env(env_vars)

Expand Down Expand Up @@ -2393,33 +2394,42 @@ def _create_config(
or container_concurrency is not None
or keep_alive_probe is not None
)
if runtime_spec is None and is_deployment_spec_updated:
has_code_or_container_spec = runtime_spec is not None
if (
mode == "create"
and not has_code_or_container_spec
and is_deployment_spec_updated
):
raise ValueError(
"To update `env_vars`, `psc_interface_config`, `min_instances`, "
"To set `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 "
"the source code options (`source_packages`, "
"`developer_connect_source`, `git_source` or `agent_config_source`)."
)

if runtime_spec is not None:
if is_deployment_spec_updated:
(
deployment_spec,
deployment_update_masks,
) = self._generate_deployment_spec_or_raise(
env_vars=env_vars,
psc_interface_config=psc_interface_config,
agent_gateway_config=agent_gateway_config,
min_instances=min_instances,
max_instances=max_instances,
resource_limits=resource_limits,
container_concurrency=container_concurrency,
keep_alive_probe=keep_alive_probe,
)
if is_deployment_spec_updated:
(
deployment_spec,
deployment_update_masks,
) = self._generate_deployment_spec_or_raise(
env_vars=env_vars,
caller_provided_env_vars=caller_provided_env_vars,
psc_interface_config=psc_interface_config,
agent_gateway_config=agent_gateway_config,
min_instances=min_instances,
max_instances=max_instances,
resource_limits=resource_limits,
container_concurrency=container_concurrency,
keep_alive_probe=keep_alive_probe,
)
if deployment_update_masks:
if runtime_spec is None:
runtime_spec = {}
update_masks.extend(deployment_update_masks)
runtime_spec["deployment_spec"] = deployment_spec

if has_code_or_container_spec:
if agent_server_mode:
if not runtime_spec.get("deployment_spec"):
runtime_spec["deployment_spec"] = (
Expand Down Expand Up @@ -2498,6 +2508,7 @@ def _generate_deployment_spec_or_raise(
self,
*,
env_vars: Optional[dict[str, Union[str, Any]]] = None,
caller_provided_env_vars: bool = True,
psc_interface_config: Optional[types.PscInterfaceConfigDict] = None,
agent_gateway_config: Optional[
types.ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict
Expand All @@ -2510,19 +2521,20 @@ def _generate_deployment_spec_or_raise(
) -> Tuple[dict[str, Any], Sequence[str]]:
deployment_spec: dict[str, Any] = {}
update_masks = []
if env_vars:
if env_vars is not None:
deployment_spec["env"] = []
deployment_spec["secret_env"] = []
if caller_provided_env_vars:
deployment_spec["secret_env"] = []
if isinstance(env_vars, dict):
self._update_deployment_spec_with_env_vars_dict_or_raise(
deployment_spec=deployment_spec,
env_vars=env_vars,
)
else:
raise TypeError(f"env_vars must be a dict, but got {type(env_vars)}.")
if deployment_spec.get("env"):
if caller_provided_env_vars or deployment_spec.get("env"):
update_masks.append("spec.deployment_spec.env")
if deployment_spec.get("secret_env"):
if caller_provided_env_vars or deployment_spec.get("secret_env"):
update_masks.append("spec.deployment_spec.secret_env")
if psc_interface_config:
deployment_spec["psc_interface_config"] = psc_interface_config
Expand Down
152 changes: 149 additions & 3 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 @@ -3252,16 +3252,162 @@ def test_update_runtime_extra_packages(

@mock.patch.object(_runtimes_utils, "_prepare")
@mock.patch.object(_runtimes_utils, "_await_operation")
def test_update_runtime_deployment_config_without_agent_raises(
def test_create_runtime_deployment_config_without_agent_raises(
self, mock_await_operation, mock_prepare
):
with pytest.raises(ValueError, match="To update `env_vars`"):
with pytest.raises(ValueError, match="To set `env_vars`"):
self.client.runtimes.create(
config=_genai_types.AgentRuntimeConfig(
env_vars=_TEST_AGENT_ENGINE_ENV_VARS_INPUT
),
)

@mock.patch.object(_runtimes_utils, "_await_operation")
def test_update_runtime_deployment_config_without_agent(
self, mock_await_operation
):
mock_await_operation.return_value = _genai_types.RuntimeOperation(
response=_genai_types.ReasoningEngine(
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
spec=_TEST_AGENT_ENGINE_SPEC,
)
)
with mock.patch.object(
self.client.runtimes._api_client, "request"
) as request_mock:
request_mock.return_value = genai_types.HttpResponse(body="")
self.client.runtimes.update(
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
config=_genai_types.AgentRuntimeConfig(
env_vars=_TEST_AGENT_ENGINE_ENV_VARS_INPUT
env_vars=_TEST_AGENT_ENGINE_ENV_VARS_INPUT,
),
)
update_mask = ",".join(
[
"spec.deployment_spec.env",
"spec.deployment_spec.secret_env",
]
)
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": {
"deployment_spec": _TEST_AGENT_ENGINE_SPEC.get(
"deployment_spec"
),
},
"_query": {"updateMask": update_mask},
},
None,
)

@mock.patch.object(_runtimes_utils, "_await_operation")
def test_update_runtime_clear_secret_env(self, mock_await_operation):
mock_await_operation.return_value = _genai_types.RuntimeOperation(
response=_genai_types.ReasoningEngine(
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
spec=_TEST_AGENT_ENGINE_SPEC,
)
)
with mock.patch.object(
self.client.runtimes._api_client, "request"
) as request_mock:
request_mock.return_value = genai_types.HttpResponse(body="")
self.client.runtimes.update(
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
config=_genai_types.AgentRuntimeConfig(
env_vars={"TEST_ENV_VAR": "TEST_ENV_VAR_VALUE"},
),
)
update_mask = ",".join(
[
"spec.deployment_spec.env",
"spec.deployment_spec.secret_env",
]
)
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": {
"deployment_spec": {
"env": [
{
"name": "TEST_ENV_VAR",
"value": "TEST_ENV_VAR_VALUE",
}
],
"secret_env": [],
},
},
"_query": {"updateMask": update_mask},
},
None,
)

@mock.patch.object(_runtimes_utils, "_await_operation")
def test_update_runtime_clear_all_env_vars(self, mock_await_operation):
mock_await_operation.return_value = _genai_types.RuntimeOperation(
response=_genai_types.ReasoningEngine(
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
spec=_TEST_AGENT_ENGINE_SPEC,
)
)
with mock.patch.object(
self.client.runtimes._api_client, "request"
) as request_mock:
request_mock.return_value = genai_types.HttpResponse(body="")
self.client.runtimes.update(
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
config=_genai_types.AgentRuntimeConfig(env_vars={}),
)
update_mask = ",".join(
[
"spec.deployment_spec.env",
"spec.deployment_spec.secret_env",
]
)
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": {
"deployment_spec": {
"env": [],
"secret_env": [],
},
},
"_query": {"updateMask": update_mask},
},
None,
)

def test_update_runtime_adk_framework_without_env_vars_preserves_env(self):
config = self.client.runtimes._create_config(
mode="update",
agent_framework="google-adk",
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
)
assert "secret_env" not in config["spec"]["deployment_spec"]
assert "spec.deployment_spec.secret_env" not in config["update_mask"]
assert "spec.deployment_spec.env" in config["update_mask"]

def test_update_runtime_empty_deployment_field_without_mask_omits_spec(
self,
):
config = self.client.runtimes._create_config(
mode="update",
psc_interface_config={},
)
assert "spec" not in config
assert "update_mask" not in config

@mock.patch.object(_runtimes_utils, "_prepare")
@mock.patch.object(_runtimes_utils, "_await_operation")
Expand Down
54 changes: 33 additions & 21 deletions vertexai/_genai/agent_engines.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2358,6 +2358,7 @@ def _create_config(
update_masks.append("traffic_config")
config["traffic_config"] = traffic_config

caller_provided_env_vars = env_vars is not None
if agent_framework == "google-adk":
env_vars = _agent_engines_utils._add_telemetry_enablement_env(env_vars)

Expand Down Expand Up @@ -2462,33 +2463,42 @@ def _create_config(
or container_concurrency is not None
or keep_alive_probe is not None
)
if agent_engine_spec is None and is_deployment_spec_updated:
has_code_or_container_spec = agent_engine_spec is not None
if (
mode == "create"
and not has_code_or_container_spec
and is_deployment_spec_updated
):
raise ValueError(
"To update `env_vars`, `psc_interface_config`, `min_instances`, "
"To set `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 "
"the source code options (`source_packages`, "
"`developer_connect_source` or `agent_config_source`)."
)

if agent_engine_spec is not None:
if is_deployment_spec_updated:
(
deployment_spec,
deployment_update_masks,
) = self._generate_deployment_spec_or_raise(
env_vars=env_vars,
psc_interface_config=psc_interface_config,
agent_gateway_config=agent_gateway_config,
min_instances=min_instances,
max_instances=max_instances,
resource_limits=resource_limits,
container_concurrency=container_concurrency,
keep_alive_probe=keep_alive_probe,
)
if is_deployment_spec_updated:
(
deployment_spec,
deployment_update_masks,
) = self._generate_deployment_spec_or_raise(
env_vars=env_vars,
caller_provided_env_vars=caller_provided_env_vars,
psc_interface_config=psc_interface_config,
agent_gateway_config=agent_gateway_config,
min_instances=min_instances,
max_instances=max_instances,
resource_limits=resource_limits,
container_concurrency=container_concurrency,
keep_alive_probe=keep_alive_probe,
)
if deployment_update_masks:
if agent_engine_spec is None:
agent_engine_spec = {}
update_masks.extend(deployment_update_masks)
agent_engine_spec["deployment_spec"] = deployment_spec

if has_code_or_container_spec:
if agent_server_mode:
if not agent_engine_spec.get("deployment_spec"):
agent_engine_spec["deployment_spec"] = (
Expand Down Expand Up @@ -2564,6 +2574,7 @@ def _generate_deployment_spec_or_raise(
self,
*,
env_vars: Optional[dict[str, Union[str, Any]]] = None,
caller_provided_env_vars: bool = True,
psc_interface_config: Optional[types.PscInterfaceConfigDict] = None,
agent_gateway_config: Optional[
types.ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict
Expand All @@ -2576,19 +2587,20 @@ def _generate_deployment_spec_or_raise(
) -> Tuple[dict[str, Any], Sequence[str]]:
deployment_spec: dict[str, Any] = {}
update_masks = []
if env_vars:
if env_vars is not None:
deployment_spec["env"] = []
deployment_spec["secret_env"] = []
if caller_provided_env_vars:
deployment_spec["secret_env"] = []
if isinstance(env_vars, dict):
self._update_deployment_spec_with_env_vars_dict_or_raise(
deployment_spec=deployment_spec,
env_vars=env_vars,
)
else:
raise TypeError(f"env_vars must be a dict, but got {type(env_vars)}.")
if deployment_spec.get("env"):
if caller_provided_env_vars or deployment_spec.get("env"):
update_masks.append("spec.deployment_spec.env")
if deployment_spec.get("secret_env"):
if caller_provided_env_vars or deployment_spec.get("secret_env"):
update_masks.append("spec.deployment_spec.secret_env")
if psc_interface_config:
deployment_spec["psc_interface_config"] = psc_interface_config
Expand Down
Loading