|
| 1 | +"""Longevity / resource-impact container test. |
| 2 | +
|
| 3 | +Exercises the XGBoost serving container under sustained inference load and |
| 4 | +tracks its resource footprint (memory + CPU) over time. The goal is to catch |
| 5 | +container-level regressions that unit / functional tests miss: |
| 6 | +
|
| 7 | +- memory leaks — RSS that climbs steadily and never plateaus under repeated |
| 8 | + ``/invocations`` requests |
| 9 | +- runaway resource use — CPU or memory that blows past a sane bound |
| 10 | +- loss of liveness — the container becoming unresponsive after prolonged load |
| 11 | +
|
| 12 | +This is intentionally lightweight (a couple of minutes) so it can run on every |
| 13 | +PR alongside the other local container tests, not just on the release pipeline. |
| 14 | +""" |
| 15 | + |
| 16 | +import logging |
| 17 | +import os |
| 18 | +import time |
| 19 | + |
| 20 | +from .container_helper import ServingContainer |
| 21 | + |
| 22 | +LOGGER = logging.getLogger(__name__) |
| 23 | + |
| 24 | + |
| 25 | +# --------------------------------------------------------------------------- |
| 26 | +# Tunables |
| 27 | +# --------------------------------------------------------------------------- |
| 28 | + |
| 29 | +# Total wall-clock duration to keep hammering the container. |
| 30 | +LOAD_DURATION_SECONDS = int(os.environ.get("XGB_LONGEVITY_DURATION", "150")) |
| 31 | +# How often to sample docker stats (memory / cpu). |
| 32 | +SAMPLE_INTERVAL_SECONDS = 5 |
| 33 | +# Warm-up window ignored when computing the memory-growth baseline — lets the |
| 34 | +# process finish lazy allocation / model load before we start trusting numbers. |
| 35 | +WARMUP_SECONDS = 20 |
| 36 | + |
| 37 | +# Memory growth from the post-warmup baseline to the end of the run that we |
| 38 | +# consider indicative of a leak. XGBoost inference on a fixed model should be |
| 39 | +# flat; a steadily climbing RSS is the signal we want to catch. |
| 40 | +MEM_GROWTH_LIMIT_MB = 150 |
| 41 | +# Absolute ceiling — the small mnist model should never need anywhere near this. |
| 42 | +MEM_ABSOLUTE_LIMIT_MB = 4096 |
| 43 | + |
| 44 | +MODEL_NAME = "mnist-xgb-model" |
| 45 | +INPUT_FILE = "mnist-1.csv" |
| 46 | +CONTENT_TYPE = "text/csv" |
| 47 | + |
| 48 | + |
| 49 | +def _model_path(resources, model_name): |
| 50 | + return os.path.join(resources, "models", model_name) |
| 51 | + |
| 52 | + |
| 53 | +def _input_path(resources, filename): |
| 54 | + return os.path.join(resources, "input", filename) |
| 55 | + |
| 56 | + |
| 57 | +def test_serving_longevity_no_memory_leak(docker_client, image_uri, inference_resources): |
| 58 | + """Sustained-load longevity check on the serving container. |
| 59 | +
|
| 60 | + Sends inference requests continuously for ``LOAD_DURATION_SECONDS`` while |
| 61 | + sampling memory and CPU. Asserts the container stays healthy, every request |
| 62 | + succeeds, and memory does not grow past the leak threshold. |
| 63 | + """ |
| 64 | + model_dir = _model_path(inference_resources, MODEL_NAME) |
| 65 | + with open(_input_path(inference_resources, INPUT_FILE), "rb") as f: |
| 66 | + payload = f.read() |
| 67 | + |
| 68 | + request_count = 0 |
| 69 | + failures = 0 |
| 70 | + samples = [] # list of (elapsed_s, mem_mb, cpu_percent) |
| 71 | + |
| 72 | + with ServingContainer(docker_client, image_uri, model_dir) as ctx: |
| 73 | + start = time.time() |
| 74 | + next_sample = 0.0 |
| 75 | + |
| 76 | + while True: |
| 77 | + elapsed = time.time() - start |
| 78 | + if elapsed >= LOAD_DURATION_SECONDS: |
| 79 | + break |
| 80 | + |
| 81 | + # Drive load. |
| 82 | + try: |
| 83 | + resp = ctx.invocations(data=payload, content_type=CONTENT_TYPE) |
| 84 | + request_count += 1 |
| 85 | + if resp.status_code != 200: |
| 86 | + failures += 1 |
| 87 | + LOGGER.warning("Non-200 response: %s %s", resp.status_code, resp.text[:200]) |
| 88 | + except Exception as exc: # connection dropped == container in trouble |
| 89 | + failures += 1 |
| 90 | + LOGGER.warning("Invocation raised: %s", exc) |
| 91 | + |
| 92 | + # Sample resources on a fixed cadence. |
| 93 | + if elapsed >= next_sample: |
| 94 | + sample = ctx.stats() |
| 95 | + samples.append((elapsed, sample["mem_mb"], sample["cpu_percent"])) |
| 96 | + LOGGER.info( |
| 97 | + "t=%5.1fs reqs=%d mem=%.1fMB cpu=%.1f%%", |
| 98 | + elapsed, |
| 99 | + request_count, |
| 100 | + sample["mem_mb"], |
| 101 | + sample["cpu_percent"], |
| 102 | + ) |
| 103 | + next_sample = elapsed + SAMPLE_INTERVAL_SECONDS |
| 104 | + |
| 105 | + # Container must still answer a health check after the load. |
| 106 | + assert ctx.ping().status_code == 200, "container unhealthy after sustained load" |
| 107 | + |
| 108 | + # -- assertions ---------------------------------------------------------- |
| 109 | + |
| 110 | + assert request_count > 0, "no requests were sent" |
| 111 | + assert samples, "no resource samples collected" |
| 112 | + assert failures == 0, f"{failures}/{request_count} invocations failed under sustained load" |
| 113 | + |
| 114 | + mem_series = [mem for _, mem, _ in samples] |
| 115 | + peak_mem = max(mem_series) |
| 116 | + cpu_series = [cpu for _, _, cpu in samples] |
| 117 | + peak_cpu = max(cpu_series) |
| 118 | + |
| 119 | + # Baseline = first sample taken after the warm-up window (fall back to the |
| 120 | + # first sample if the run is shorter than warm-up for any reason). |
| 121 | + post_warmup = [mem for elapsed, mem, _ in samples if elapsed >= WARMUP_SECONDS] |
| 122 | + baseline_mem = post_warmup[0] if post_warmup else mem_series[0] |
| 123 | + final_mem = mem_series[-1] |
| 124 | + mem_growth = final_mem - baseline_mem |
| 125 | + |
| 126 | + LOGGER.info( |
| 127 | + "longevity summary: requests=%d duration=%ds peak_mem=%.1fMB " |
| 128 | + "baseline_mem=%.1fMB final_mem=%.1fMB growth=%.1fMB peak_cpu=%.1f%%", |
| 129 | + request_count, |
| 130 | + LOAD_DURATION_SECONDS, |
| 131 | + peak_mem, |
| 132 | + baseline_mem, |
| 133 | + final_mem, |
| 134 | + mem_growth, |
| 135 | + peak_cpu, |
| 136 | + ) |
| 137 | + |
| 138 | + assert peak_mem < MEM_ABSOLUTE_LIMIT_MB, ( |
| 139 | + f"peak memory {peak_mem:.1f}MB exceeded absolute ceiling {MEM_ABSOLUTE_LIMIT_MB}MB" |
| 140 | + ) |
| 141 | + assert mem_growth < MEM_GROWTH_LIMIT_MB, ( |
| 142 | + f"memory grew {mem_growth:.1f}MB (baseline {baseline_mem:.1f}MB -> " |
| 143 | + f"final {final_mem:.1f}MB) over {LOAD_DURATION_SECONDS}s, exceeding leak " |
| 144 | + f"threshold {MEM_GROWTH_LIMIT_MB}MB" |
| 145 | + ) |
0 commit comments