diff --git a/agentplatform/_genai/_runtimes_utils.py b/agentplatform/_genai/_runtimes_utils.py index b037caebb8..a151adc3ae 100644 --- a/agentplatform/_genai/_runtimes_utils.py +++ b/agentplatform/_genai/_runtimes_utils.py @@ -1294,11 +1294,6 @@ def _upload_runtime( "`from a2a import types`." ) raise TypeError(error_msg) from e - with blob.open("rb") as f: - try: - _ = cloudpickle.load(f) - except Exception as e: - raise TypeError("agent runtime serialized to an invalid format") from e dir_name = f"gs://{gcs_bucket.name}/{gcs_dir_name}" logger.info(f"Wrote to {dir_name}/{_BLOB_FILENAME}") diff --git a/tests/unit/agentplatform/genai/replays/test_create_agent_engine_a2a.py b/tests/unit/agentplatform/genai/replays/test_create_agent_engine_a2a.py index 015736d8e3..abdb76f17e 100644 --- a/tests/unit/agentplatform/genai/replays/test_create_agent_engine_a2a.py +++ b/tests/unit/agentplatform/genai/replays/test_create_agent_engine_a2a.py @@ -14,9 +14,7 @@ # # pylint: disable=protected-access,bad-continuation,missing-function-docstring -import contextlib import os -from unittest import mock from tests.unit.agentplatform.genai.replays import pytest_helper from agentplatform._genai import types from agentplatform.frameworks.a2a import default_a2a_agent @@ -30,7 +28,7 @@ from a2a.utils.constants import TransportProtocol -def test_create_a2a_agent(client, is_replay_mode): +def test_create_a2a_agent(client): my_agent = default_a2a_agent() my_agent.agent_card.supported_interfaces.append( @@ -44,30 +42,20 @@ def test_create_a2a_agent(client, is_replay_mode): staging_bucket = os.environ["GCS_BUCKET"] - # In replay mode, GCS operations are mocked and blob.open("rb") returns a mock - # that fails when cloudpickle.load expects bytes. We mock _upload_runtime - # to skip this verification step, which is not needed when replaying API calls. - upload_patch = ( - mock.patch("agentplatform._genai._runtimes_utils._upload_runtime") - if is_replay_mode - else contextlib.nullcontext() + runtime = client.runtimes.create( + agent=my_agent, + config={ + "staging_bucket": staging_bucket, + "display_name": "test-a2a-agent", + "http_options": {"api_version": "v1beta1"}, + "requirements": [ + "google-cloud-aiplatform[agent_engines]", + "a2a-sdk", + "sse-starlette", + ], + }, ) - with upload_patch: - runtime = client.runtimes.create( - agent=my_agent, - config={ - "staging_bucket": staging_bucket, - "display_name": "test-a2a-agent", - "http_options": {"api_version": "v1beta1"}, - "requirements": [ - "google-cloud-aiplatform[agent_engines]", - "a2a-sdk", - "sse-starlette", - ], - }, - ) - assert isinstance(runtime, types.Runtime) assert runtime.api_resource.display_name == "test-a2a-agent" diff --git a/vertexai/_genai/_agent_engines_utils.py b/vertexai/_genai/_agent_engines_utils.py index 2da175b745..08ca82b652 100644 --- a/vertexai/_genai/_agent_engines_utils.py +++ b/vertexai/_genai/_agent_engines_utils.py @@ -1281,11 +1281,6 @@ def _upload_agent_engine( "`from a2a import types`." ) raise TypeError(error_msg) from e - with blob.open("rb") as f: - try: - _ = cloudpickle.load(f) - except Exception as e: - raise TypeError("Agent engine serialized to an invalid format") from e dir_name = f"gs://{gcs_bucket.name}/{gcs_dir_name}" logger.info(f"Wrote to {dir_name}/{_BLOB_FILENAME}") diff --git a/vertexai/agent_engines/_agent_engines.py b/vertexai/agent_engines/_agent_engines.py index fa51907dd6..eb96078e04 100644 --- a/vertexai/agent_engines/_agent_engines.py +++ b/vertexai/agent_engines/_agent_engines.py @@ -1302,11 +1302,6 @@ def _upload_agent_engine( "`from a2a import types as a2a_types`." ) raise TypeError(error_msg) from e - with blob.open("rb") as f: - try: - _ = cloudpickle.load(f) - except Exception as e: - raise TypeError("Agent engine serialized to an invalid format") from e dir_name = f"gs://{gcs_bucket.name}/{gcs_dir_name}" logger.info(f"Wrote to {dir_name}/{_BLOB_FILENAME}")