add agentic benchmarking on gke - #6772
Conversation
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
| # Used with --gke_provision_mode=native | ||
| # | ||
| # Prerequisites (run once before PKB): | ||
| # python tools/agentic-benchmark/scripts/prerequisite_setup.py \ |
There was a problem hiding this comment.
This tool isn't being included & therefore this comment doesn't need to be here.
There was a problem hiding this comment.
Done, removed stale references in new commit.
| # For sweeps (cluster pre-exists, PKB skips provision/teardown): | ||
| # The sweep bridge injects --run_stage=run,cleanup automatically. | ||
|
|
||
| gke_python_density: |
There was a problem hiding this comment.
Internally we put a lot of this info but externally it is useful.. it's probably a good addition.
| @@ -0,0 +1,240 @@ | |||
| from google.adk.agents import LlmAgent | |||
| from google.adk.code_executors import GkeCodeExecutor | |||
There was a problem hiding this comment.
Where is this file run? From the same machine running PKB or a different one?
There was a problem hiding this comment.
This is the ADK Agent we're benchmarking against (i.e calling its FASTAPI APIs). It Get's Docker-built and deployed to GKE. PKB Benchmarks target it via kubectl port-forward.
| six>=1.13.0 | ||
| timeout-decorator | ||
| scipy | ||
| matplotlib |
There was a problem hiding this comment.
I don't see a reference to this elsewhere with a ctrl-f; is it leftover from an earlier version?
In general we prefer not making many changes to requirements.txt.
There was a problem hiding this comment.
Stale, Removed.
| ' beyond the default node pool (e.g. kubernetes_node_scale with 5k nodes).', | ||
| ) | ||
|
|
||
| GKE_USE_BETA = flags.DEFINE_boolean( |
There was a problem hiding this comment.
If we add this flag, IMO just make it "gcloud_use_beta" (or actually an enum use alpha, beta, None "gcloud_beta_version") & being referenced from gcp/util.py directly seems best.
Alternatively we often will say in the provider "if preview feature used, cmd.use_beta_gcloud = True". In general what feature are you using that needs beta?
There was a problem hiding this comment.
I removed this Flag now. It was for --enable-pod-snapshots on GKE, which was BETA at the time of development.
| by all seven UC benchmark scripts. Each benchmark's Provision() and | ||
| Teardown() functions delegate to the public functions in this module. | ||
|
|
||
| Infrastructure created (in order): |
There was a problem hiding this comment.
The very premise of this file is incorrect. PKB (and esp eg google_kubernetes_engine.py _Create) should be handling all of the provisioning logic.
I'm not sure how much of this is a) completely unnecessary because it's handled elsewhere in PKB (like we do setup subnets & networks automatically if you don't specify a network" or b) is indeed necessary but should be located in some other Resource.py class.
There was a problem hiding this comment.
+1. Let's set up the cloud infra using PKB-native way.
There was a problem hiding this comment.
There were 2 approaches for Provisioning when this comment was made; 'Custom', and PKB 'Native'. I removed the 'Custom' option and any unnecessary code related to Custom; PKB-Native is the only way now.
| chromium_replicas = FLAGS.gke_chromium_replicas | ||
|
|
||
| manifest = """--- | ||
| apiVersion: extensions.agents.x-k8s.io/v1alpha1 |
There was a problem hiding this comment.
should go in some .yaml.j2 file
There was a problem hiding this comment.
Done, moved all inline-templates to Jinja2 templates.
| return _RunCmd(cmd, check=check, timeout=timeout) | ||
|
|
||
|
|
||
| def _KubectlApply(manifest_str): |
There was a problem hiding this comment.
why have you rewritten kubectl apply & _RunKubectl when implementations exist container_service/kubectl.py ?
There was a problem hiding this comment.
Done, moved refactored all to use kubectl.py
| @@ -0,0 +1,362 @@ | |||
| """PKB Benchmark: GKE Agent Python Sandbox Density (Use Case B). | |||
There was a problem hiding this comment.
For easier review and faster iteration, I'd recommend keeping one benchmark in this PR and leave the other benchmarks for followup PRs. My recommendation is to keep the Python density benchmark.
| @@ -0,0 +1,362 @@ | |||
| """PKB Benchmark: GKE Agent Python Sandbox Density (Use Case B). | |||
There was a problem hiding this comment.
Let's drop "(Use Case B)" from the description. For the published PKB benchmarks, the documentation should clearly state what the benchmarks are about. The ordering of A,B,C... will become stale and confusing to readers.
| @@ -0,0 +1,362 @@ | |||
| """PKB Benchmark: GKE Agent Python Sandbox Density (Use Case B). | |||
There was a problem hiding this comment.
Can we drop "GKE" from the file name and the description? Based on the path this is a Kubernetes benchmark, and presumably this benchmark can be reused for other cloud provider without significant change, right?
There was a problem hiding this comment.
The benchmarks have GKE-specific dependencies at the moment, such as Pod Snapshots (podsnapshot.gke.io/v1 CRD), image building using cloud build, ... etc. Abstracting this coupling would require some research and refactoring, and possibly another PR.
| # --------------------------------------------------------------------------- | ||
|
|
||
| flags.DEFINE_integer( | ||
| "gke_python_density", |
There was a problem hiding this comment.
gke_python_density
nit: Shall we name the flag something like "concurrent_sandbox_count"?gkeandpythoncan already be implied based on the file name and description of the benchmark.
There was a problem hiding this comment.
renamed to 'gke_python_density_concurrent_sandbox_count'.
Benchmark-/usecase- specific flags should maintain the benchmark name (for example: gke_python_density_) as a prefix in order not to have any potential 'namespace collisions' when multiple benchmarks are imported or executed.
| flags.DEFINE_integer( | ||
| "gke_python_density_sample_warmup", | ||
| 0, | ||
| "Number of warmup iterations per session (excluded from stats).", |
There was a problem hiding this comment.
It's unclear what "warmup iterations" means as it's not mentioned before. Shall we document the workflow in the benchmark description?
There was a problem hiding this comment.
added a description to the docstring at the top of the file.
| by all seven UC benchmark scripts. Each benchmark's Provision() and | ||
| Teardown() functions delegate to the public functions in this module. | ||
|
|
||
| Infrastructure created (in order): |
There was a problem hiding this comment.
+1. Let's set up the cloud infra using PKB-native way.
| # --------------------------------------------------------------------------- | ||
|
|
||
|
|
||
| def _emit(samples, agg, agg_key, metric_suffix, unit, namespace, extra): |
There was a problem hiding this comment.
Can you document how the metrics emit works and what the parameters are?
| # --------------------------------------------------------------------------- | ||
|
|
||
|
|
||
| def _BuildADKAgentImage( |
There was a problem hiding this comment.
Probably this whole gke_image_build_utils file is not needed. See GoogleArtifactRegistry in google_kubernetes_engine.py & kubernetes_hpa
with
container_specs:
kubernetes_fib:
image: fibonacci
in the config + data/docker/fibonacci for the dockerimage
There was a problem hiding this comment.
We're building 3 images.
The ADK Image, whose codebase is in the repo, and two other images (Sandbox Router and Chrom Sandbox) whose codebase is not in the repo; they're in https://github.com/kubernetes-sigs/agent-sandbox, and they are not built/published publically, they need to be built per use.
I can try to move ADK Image to PKB NAtive, but the other 2 will still require to be in a "repreq" python script, and that is because I do not want to import their static code into PKB.
| # Chrome and Router images are built during prerequisites | ||
| # (gke_prerequisites.py), not during Prepare. | ||
| # ADK agent image is built by PKB container_specs during Provision. | ||
| from perfkitbenchmarker.linux_benchmarks.kubernetes.agentic import ( |
There was a problem hiding this comment.
across the board avoid importing within functions & import at the top instead.
| 4. Bind IAM roles | ||
| 5. Deploy PodSnapshotStorageConfig + PodSnapshotPolicy | ||
| """ | ||
| if FLAGS.skip_deploy_snapshots: |
There was a problem hiding this comment.
Frequently prefer using https://absl.readthedocs.io/en/latest/absl.flags.html#absl.flags.FlagHolder which shows where the flag is defined.
| ) | ||
| ) | ||
|
|
||
| if n > 0: |
There was a problem hiding this comment.
Consider throwing if eg n == 0 would be consider an error. In general throw somewhat liberally rather than handling errors that resulted in unexpected state.
| # --------------------------------------------------------------------------- | ||
|
|
||
| flags.DEFINE_string( | ||
| "k8s_namespace", |
There was a problem hiding this comment.
should be named k8s_agentic_namespace & could possibly just be a hardcoded constant K8S_AGENTIC_NAMESPACE = "agentic" as well.
| ) | ||
|
|
||
| flags.DEFINE_bool( | ||
| "k8s_gvisor", |
There was a problem hiding this comment.
likewise prefix all the flags with k8s_agent since they're only used in these files.
| ) | ||
|
|
||
|
|
||
| def CountPods(namespace, label, phase=None): |
There was a problem hiding this comment.
Use pytyping pretty much everywhere: https://google.github.io/pytype/user_guide.html
| import atexit | ||
| import os as _os | ||
| import signal | ||
| import threading |
There was a problem hiding this comment.
same comments about imports at the top. Probably just move this to its own file given its independence
| # Derive machine_type from benchmark_spec (set via set_benchmark_spec) | ||
| machine_type = None | ||
| if _current_benchmark_spec: | ||
| cluster = getattr(_current_benchmark_spec, 'container_cluster', None) |
There was a problem hiding this comment.
is this 2 or 4 spaces for the indent? Should be 2. That might be handled by the linter from CONTRIBUTING.md or might not be.
| def set_benchmark_spec(benchmark_spec): | ||
| """Store benchmark_spec for metadata derivation (called by Run()).""" | ||
| global _current_benchmark_spec | ||
| _current_benchmark_spec = benchmark_spec |
There was a problem hiding this comment.
This is kind of a smell. What it indicates to me is that you should be considering using more classes to manage state. In general IMO there ought to be some "Sandbox" class/resource. A resource is expected to handle creation,deletion, metadata, & flags/spec. I believe you are doing all of those things just with global functions.
See the new README in https://github.com/GoogleCloudPlatform/PerfKitBenchmarker/tree/master/perfkitbenchmarker/resources
There was a problem hiding this comment.
there seems to be work done on that in different PRs.
https://github.com/GoogleCloudPlatform/PerfKitBenchmarker/pull/6732/changes
https://github.com/GoogleCloudPlatform/PerfKitBenchmarker/pull/6740/changes
We can just refactor and consume those resources in a different Effort.
| -> GkeCodeExecutor -> SandboxClient -> gVisor sandbox pod | ||
| """ | ||
|
|
||
| """GKE Performance Agent â ADK agent definition for sandbox benchmarking. |
| with open(payload_script_path, "r") as f: | ||
| payload_benchmark_code = f.read() | ||
| except Exception: | ||
| payload_benchmark_code = "import os; print(os.uname())" |
There was a problem hiding this comment.
This is a programming error. It should fail more loudly rather than continuing in an invalid state.
There was a problem hiding this comment.
Agree. Will put in a fix to crash the agent pod (CrashLoopBackOff) when reading the file fails.
| with open(density_script_path, "r") as f: | ||
| density_benchmark_code = f.read() | ||
| except Exception: | ||
| density_benchmark_code = "import os; print(os.uname())" |
There was a problem hiding this comment.
This is a programming error. It should fail more loudly rather than continuing in an invalid state.
There was a problem hiding this comment.
Agree. Will put in a fix to crash the agent pod (CrashLoopBackOff) when reading the file fails.
| with open(qps_script_path, "r") as f: | ||
| qps_benchmark_code = f.read() | ||
| except Exception: | ||
| qps_benchmark_code = "import json; print(json.dumps({'sandbox_status': 'ok'}))" |
There was a problem hiding this comment.
This is a programming error. It should fail more loudly rather than continuing in an invalid state.
There was a problem hiding this comment.
Will Fix.Agree. Will put in a fix to crash the agent pod (CrashLoopBackOff) when reading the file fails.
| with open(qps_script_path, "r") as f: | ||
| qps_benchmark_code = f.read() | ||
| except Exception: | ||
| qps_benchmark_code = "import json; print(json.dumps({'sandbox_status': 'ok'}))" |
There was a problem hiding this comment.
Write a reusable function for the scripts loading:
density_benchmark_code = _load_script("benchmark_density.py")
payload_benchmark_code = _load_script("benchmark_payload.py")
qps_benchmark_code = _load_script("benchmark_qps.py")
| # Module-level thread pool for sandbox I/O operations. | ||
| # Initialized once at import time to avoid thread-safety issues | ||
| # with lazy initialization inside _execute_in_sandbox(). | ||
| _SANDBOX_POOL = ThreadPoolExecutor(max_workers=16) |
There was a problem hiding this comment.
The number of workers should be configurable benchmark config. Otherwise we are limited with 16 concurrent agent sandbox executions.
| from k8s_agent_sandbox.sandbox_client import SandboxClient | ||
| from k8s_agent_sandbox.models import SandboxDirectConnectionConfig | ||
| import logging | ||
| import time |
There was a problem hiding this comment.
As we discussed in the meeting, please move imports to the top of the file.
| _SANDBOX_POOL = ThreadPoolExecutor(max_workers=16) | ||
|
|
||
|
|
||
| class V3GkeCodeExecutor(GkeCodeExecutor): |
There was a problem hiding this comment.
just a lingering development iteration name from iterating against older k8s-agent-sandbox API versions. Will rename to BenchmarkGkeCodeExecutor, more meaningful.
|
|
||
| class V3GkeCodeExecutor(GkeCodeExecutor): | ||
| def _execute_in_sandbox(self, code: str) -> CodeExecutionResult: | ||
| """Executes code using the v0.4.6 compatible SandboxClient.""" |
There was a problem hiding this comment.
Let's drop the version "v0.4.6" from the code, since it will easily get stale. You already have the version tracked in requirements.txt
| @@ -0,0 +1,276 @@ | |||
| """GKE Performance Agent -- ADK agent definition. | |||
There was a problem hiding this comment.
There is no GKE-specific content in this file. Let's replace "GKE" with "k8s"
| # Measure TTFE | ||
| ttfe_start = time.perf_counter() | ||
| exec("x = 1 + 1", globals()) | ||
| results["ttfe_ms"] = round((time.perf_counter() - ttfe_start) * 1000, 6) |
There was a problem hiding this comment.
This is wrong. Time-To-First-Execution is the time between the agent triggering execution and the sandbox conducting the actual execution. Here you're measuring a period of time within the sandbox
There was a problem hiding this comment.
Will remove, this is within sandbox anyways. the real TTFE is orchestrator-side measured elsewhere in code.
|
|
||
| tasks = get_static_tasks() | ||
| sampled_tasks = [t for t in tasks if t["type"] != "import"] | ||
| import_task = next((t for t in tasks if t["type"] == "import"), None) |
There was a problem hiding this comment.
It's weird that we are using next() while the hardcoded list has exactly one "import" task. Why is that?
There was a problem hiding this comment.
because import_task = tasks[2] is fragile.
If next is unapealing, I can replace with import_task = [t for t in tasks if t["type"] == "import"][0], same effect.
| # error on repeated reimport, so it runs once outside the loop) | ||
| for _ in range(SAMPLE_WARMUP): | ||
| for task in sampled_tasks: | ||
| exec(task["code"], globals()) |
There was a problem hiding this comment.
Can you document why we need a warmup phase here? What things are we warming up?
There was a problem hiding this comment.
will add documentation, something like:
# Warmup iterations stabilize measurements by pre-warming:
# - Python bytecode cache (first exec compiles .pyc)
# - gVisor sentry page cache (first memory access triggers page faults)
# - OS dentry/inode cache (first stat/listdir populates kernel caches)
# - CPU branch predictor (first iterations train prediction tables)
| ] | ||
|
|
||
|
|
||
| def _percentile(sorted_vals, pct): |
There was a problem hiding this comment.
_percentile is not used in this file. Drop it?
There was a problem hiding this comment.
will remove.
|
|
||
| print(json.dumps({ | ||
| "sandbox_status": "ok", | ||
| "sandbox_qps_exec_ms": round(elapsed_ms, 3), |
There was a problem hiding this comment.
Let's drop "qps" from the metrics name since this metrics measures how long a simple python sandbox execution takes-- nothing related to qps here.
| print(json.dumps({ | ||
| "sandbox_status": "ok", | ||
| "sandbox_qps_exec_ms": round(elapsed_ms, 3), | ||
| "sandbox_compute_result": result, |
There was a problem hiding this comment.
Let's drop the "sandbox_compute_result" since it's just confusing to people unless they read the actual source of what the sandbox does.
| elapsed_ms = (time.perf_counter() - t0) * 1000 | ||
|
|
||
| print(json.dumps({ | ||
| "sandbox_status": "ok", |
There was a problem hiding this comment.
Let's make this a "success" v.s. "fail" based on whether result matches the expectation (49,995,000-- the sum of all integers from 0 up to 9,999)
| @@ -0,0 +1,199 @@ | |||
| #!/usr/bin/env python3 | |||
| """Agentic Payload Transfer Benchmark (Use Case D). | |||
|
|
||
| def get_rss_mb() -> float: | ||
| """Get current RSS memory in MB.""" | ||
| return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024 |
There was a problem hiding this comment.
This function is also defined in benchmark_density.py. Can we have it in a util package and reuse it?
| # --- Warmup (2 iterations, not recorded) --- | ||
| # Stabilizes os.urandom() entropy pool and base64 codec | ||
| # initialization before measured iterations begin. | ||
| for _ in range(2): |
There was a problem hiding this comment.
Similar to the density benchmark, can you document what the warmup phase is for, and make the number of iterations configurable?
| POST /benchmark/chromium/density → run the Chromium density benchmark (UC-C) | ||
| POST /run → raw ADK agent interaction | ||
|
|
||
| POST /benchmark/python/density — Request: |
There was a problem hiding this comment.
The documentation seems incomplete as it only has the request and response for POST /benchmark/python/density. I'd recommend putting the API endpoint documentation in a separate doc file.
| logging.exception("Error shutting down ThreadPoolExecutor") | ||
|
|
||
|
|
||
| app = FastAPI(title="GKE Benchmark Agent", version="0.2.0", lifespan=lifespan) |
| _benchmark_lock = asyncio.Lock() | ||
|
|
||
|
|
||
| def _percentile_stats(sorted_values: list, prefix: str) -> dict: |
There was a problem hiding this comment.
Can you turn this function into a util for reuse?
| ) | ||
|
|
||
|
|
||
| class ChromiumBenchmarkRequest(BaseModel): |
There was a problem hiding this comment.
Right now the main.py combines all the endpoints and their Request class. For better readability, can you split each endpoint and its Request class into a separate file?
| req.concurrent_sessions, | ||
| ) | ||
|
|
||
| prompt = "Please start the GKE performance benchmark workflow." |
There was a problem hiding this comment.
The prompt is not used, correct? Can we drop it?
| sandbox-side metrics. | ||
| """ | ||
| async with _benchmark_lock: | ||
| os.environ["BENCHMARK_MODE"] = "density" |
There was a problem hiding this comment.
Not urgent, but maybe for a followup PR:
Using os env to pass configs is not very intuitive. Maybe consider using the prompt to pass the configs
POST /benchmark/python/density
|
v
FastAPI service
|
| (prompt: {mode: "density", sample_count: 100...})
v
ADK agent
|
| (prompt: {mode: "density", sample_count: 100...})
v
MockLLM
|
| (script instruction based on the prompt)
v
ADK agent (execution)
e2630db to
60f7070
Compare
3 participants
Agentic Workload Benchmarking for GKE (PKB Extension)
Summary
Adds a complete benchmarking framework for Agentic Workloads on Google Kubernetes Engine (GKE) — specifically measuring per-operation performance of untrusted Python code execution and headless Chromium browser tasks running under gVisor (GKE Agent Sandbox) isolation.
Motivation
AI agent systems require ephemeral, isolated execution environments (sandboxes) for running untrusted code. Understanding the performance characteristics of these sandboxes under gVisor — including cold-start latency, execution overhead, memory density limits, and scheduling throughput — is critical for production capacity planning.
This framework enables systematic, repeatable measurement of these characteristics across multiple GCP machine families.
Architecture
Benchmark Definitions (7 Use Cases)
gke_snapshotgke_python_densitygke_chromium_densitygke_payloadgke_warmpoolgke_qpsgke_deletionShared Utilities
gke_benchmark_utils.pygke_deploy_utils.pygke_provision_utils.pygke_image_build_utils.pygke_prerequisite_setup.pyDual Provisioning Modes
custommode: Directgcloudcalls for full infrastructure controlnativemode: Uses PKB's built-incontainer_clusterprovisioner with prerequisite script for resources PKB cannot managePKB Provider Extensions
Small additions to support GKE preview features:
--gke_use_betaflag (forcesgcloud beta container clusters create)--gke_additional_flagslist (appended to cluster create)--gke_additional_nodepool_flagslist (appended to node pool create)In-Cluster Components
ADK Agent (
workloads/adk_agent/)A FastAPI service deployed inside GKE that:
/benchmark/python/density,/benchmark/python/payload,/benchmark/python/qps,/benchmark/chromium/density)DirectConnection(in-cluster) orkubectl port-forward(dev mode)Sandbox Scripts (
sandboxed_apps/)benchmark_density.py— CPU-bound, syscall-heavy, and import-heavy tasks with RSS trackingbenchmark_payload.py— Payload generation, serialization, and stdout transfer measurementbenchmark_qps.py— Minimal script proving sandbox livenessbenchmark_density.js— Playwright-driven Chromium interaction benchmarkVibe Coding Workloads (
workloads/vibe_coding/)Startup scripts simulating real-world agentic cold-starts:
startup_pip_fastapi.sh— pip install + FastAPI server bootstartup_npm_vite.sh— npm install + Vite dev server bootUsage
Prerequisites (once per environment)
python -m perfkitbenchmarker.linux_benchmarks.kubernetes.agentic.gke_prerequisite_setup \ --project_id=sada-gke-benchmarking2 \ --region=us-central1 \ --zone=us-central1-a \ --machine_type=c4-standard-8Provision Cluster
python pkb.py --benchmarks=gke_python_density \ --run_stage=provision \ --gke_provision_mode=native \ --project=sada-gke-benchmarking2 \ --owner=george-kalisse \ --benchmark_config_file=k8s_agents/config/native_provision_config.yaml \ --gce_network_name=george-agentic-vpc \ --gce_subnet_region=us-central1 \ --zone=us-central1-a \ --container_cluster_version=1.35.3-gke.1389000 \ --gke_use_beta=true \ --gke_additional_flags="--enable-pod-snapshots,--enable-dataplane-v2,--enable-private-nodes,--enable-ip-alias,--master-ipv4-cidr=172.16.0.0/28,--workload-pool=sada-gke-benchmarking2.svc.id.goog,--subnetwork=george-agentic-subnet,--enable-master-authorized-networks,--master-authorized-networks=$(curl -s ifconfig.me)/32" \ --gke_additional_nodepool_flags="--max-pods-per-node=250" \ --gke_enable_shielded_nodes=false \ --run_uri=test \ --temp_dir=./testing/pkb/c4-standard-8/ucb``` ### Run Benchmark ```bash python pkb.py --benchmarks=gke_python_density \ --run_stage=prepare,run,cleanup \ --gke_provision_mode=native \ --gke_project_id=sada-gke-benchmarking2 \ --gke_region=us-central1 \ --gke_zone=us-central1-a \ --gke_sandbox_machine_type=c4-standard-8 \ --gke_namespace=agentic \ --gke_sandbox_version=v0.4.6 \ --gke_python_density=4 \ --gke_python_density_sample_count=20 \ --gke_python_density_sample_warmup=0 \ --gke_python_density_patch_warmpool=true \ --gke_python_density_exec_timeout=600 \ --gke_machine_type=c4-standard-8 \ --gke_gvisor=true \ --gke_api_url=http://localhost:8080 \ --run_uri=test \ --temp_dir=./testing/pkb/c4-standard-8/ucb