robust mproxy sidecar and billing-grade GCS cost tracking#1228
Merged
robust mproxy sidecar and billing-grade GCS cost tracking#1228
Conversation
Move env-var setup, CA volume, wait_for_ca script, and the two sidecar builders out of get_pod_spec. Dedupes the V1Container boilerplate via _build_sidecar_container. No behavior change.
mproxy: cpu=200m/mem=512Mi requests, 1Gi memory limit, TCP startup+ liveness probes on :8080. oom_tracker: cpu=50m/mem=128Mi requests, 256Mi limit. Without guaranteed resources mproxy was starved during main's CPU bursts and its TLS handshakes deadlocked — ProxyError on main's GCS calls. Probes also catch hangs, not just crashes.
_check_pod_disruption walks all container statuses and tags messages with container=<name>, so an mproxy OOM is distinguishable from a main OOM. watch_for_run_events widens the "container main" filter to also allow "container mproxy" so its K8s events land in events.log.
_resilient_watch moves the try/except inside the while loop with capped backoff, so a transient GKE-master connection error no longer permanently disables disruption detection. Both watchers collapse to thin adapters on top of shared _resilient_watch / _open_watcher_log / _append_watcher_log helpers.
wait_for_ca bounds its wait at 120s and exits 1 on timeout instead of falling back to a "tracking disabled" path. _setup_mitmproxy_ca raises on any setup failure; run_sidecar no longer has the "mproxy unavailable" branch. Dockerfile asserts mitmdump is installed. GCS cost tracking is a billing requirement — a broken mproxy must manifest as a visible pod-level crash-loop, not silently run without tracking.
Extract pure-Python classifier + GCSStats into gcs_classification (no zetta_utils deps) so the mitmproxy addon subprocess can use them without pulling in firestore. Removes the duplicate copies in mitm_addon that had drifted. Tightens the addon's host filter to an exact match.
extract_bucket_from_api_path now recognises XML-style /{bucket}/{object}
paths (cloudvolume / tensorstore use these — they were the bulk of the
_unknown bucket attribution). The classifier's JSON-API substring
heuristics for /o and /b are gated on the path actually being JSON,
fixing the regression where XML object reads with /o substrings
(.../object, .../output) flipped from Class B get to Class A list_objects.
17 new unit tests cover XML cases and the inflation regression.
extract_bucket_from_api_path returning None for a non-batch path is a classifier coverage gap, not a real bucket. Record under _unclassified with op_class=None (operations count + egress increment, billed totals do not), and log the full method+URL at WARNING so the gap can be fixed. GCSStats.record now accepts op_class=None for this and the upcoming _batch case.
route_request now categorises every request into one of: real bucket (classified), _batch (POST /batch/storage/v1; sub-ops not yet parsed, so uncounted), or _unclassified (coverage gap; callback logs the URL). mitm_addon._record collapses to a one-line wrapper. _batch and _unclassified make billing-impacting categories visible without inflating Class A totals.
Move classifier / extractor / GCSStats / route_request tests out of test_gcs_tracker.py into a new test_gcs_classifier.py, parametrize the per-row cases, and add a captured-traffic fixture asserting expected (bucket, class, operation) end-to-end via route_request. A coverage guard fails if any fixture row falls through to _unclassified, so a classifier-coverage gap surfaces as a test failure rather than silent cost drift.
Egress was only counted on GET responses and used the decoded body length. GCP bills wire bytes for egress, and listings / metadata / resumable-upload status responses also have egress. Now counts on any response with a body, using Content-Length when present (with len(body) fallback for chunked / multipart).
Procedure to diff per-run sidecar totals (Class A/B ops, egress) against GCP BigQuery billing export on a dedicated project. Run after classifier changes, before promising new tracked SKUs to a customer, and on a quarterly cadence. Captures per-discrepancy iteration history.
- Resource reservations: mproxy 1Gi/2Gi, oom_tracker 256Mi/512Mi (previous limits OOM'd under sustained concurrent TLS + firestore imports). - stream_large_bodies=64k: stream large response bodies instead of buffering them for the addon. - Self-heal: _run_main_loop restarts mitmdump in-place (bounded) when it crashes, so port 8080 stays bound without depending on K8s restart-policy. - Drop TCP startup probe, relax liveness to 60s grace — self-heal + wait_for_ca already cover readiness; the probe was actively harmful under chaos. - Mazepa worker pods use restart_policy=OnFailure so kubelet restarts a dead sidecar in-place. Training keeps Never (pod-level recovery via torch elastic).
- Preserve stats across mitmdump restart: mitm_addon resumes from /tmp/gcs_tracker_stats.json on startup so self-heal does not reset _stats to empty. Collapse module-level side effects into _init(). - atexit flush in gcs_tracker alongside the SIGTERM handler so any clean Python exit path flushes stats to Firestore. - Count egress on streamed bodies: drop the "if body else 0" guard so compute_egress_bytes falls back to Content-Length when stream_large_bodies is in effect. Observed 320 MB tracked on a workload that previously reported 2.5 MB. - Per-bucket egress_bytes zeroed for same-region buckets at aggregation, matching the cross-region-only top-level total. - Treat "Unable to connect to proxy" as a transient error so GCS calls that hit mproxy during a restart window retry via SQS visibility timeout instead of aborting the run.
CloudFiles-driven flow that exercises every classifier branch in one task: XML puts, gets, heads, JSON list, batch delete. Used by the stress spec to validate stats preservation and egress accuracy without a heavy real workload.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #1228 +/- ##
=======================================
Coverage 99.98% 99.98%
=======================================
Files 200 200
Lines 10642 10642
=======================================
Hits 10640 10640
Misses 2 2 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
dodamih
approved these changes
Apr 21, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Sidecar robustness
_run_main_looprestarts mitmdump in-place when it crashes.stream_large_bodies=64kkeeps mproxy memory bounded.restart_policy=OnFailurefor mazepa workers so kubelet restarts a dead sidecar in-place (training keepsNever).wait_for_cafails non-zero on timeout.container=<name>, and covers mproxy events.Cost tracking accuracy
gcs_classification.py(stdlib-only) used by both sidecar and addon./{bucket}/{object}) now handled correctly — root cause of the 2× inflation._batch(uncounted) instead of_unknown._unclassifiedis a bug signal — unrecognised paths log loudly, don't silently fold into Class A.Content-Length(was missing: 320 MB vs 2.5 MB observed).atexitflush).