Amazon Bedrock AgentCore Runtime V2 の課金メモリがピーク後に縮小する挙動を USAGE_LOGS で観測してみた

"""
AgentCore Runtime V2 課金メモリ観測エージェント

/invocations に渡せるフィールド:
  alloc_mb    : ヒープ確保量 (MiB, default=0)
  sleep_sec   : 解放後にスリープする秒数 (default=0)
  interval_sec: sleep_sec > 0 の場合、interval_sec ごとに VmRSS と memory.current を記録 (default=10)

返却フィールド:
  vmrss          : /proc/self/status の VmRSS 増分（確保前後・解放後）
  cgroup_memory  : /sys/fs/cgroup/memory.current (存在しない場合は null)
  sleep_snapshots: sleep_sec > 0 の場合、interval_sec ごとの VmRSS と memory.current のスナップショット列
"""
import gc
import json
import os
import time
from flask import Flask, request, jsonify, Response

app = Flask(__name__)

CGROUP_MEMORY_CURRENT = "/sys/fs/cgroup/memory.current"
CGROUP_MEMORY_STAT    = "/sys/fs/cgroup/memory.stat"

def read_proc_status():
    """
    /proc/self/status からメモリ関連行を dict で返す。
    単位は kB のまま返す。
    """
    mem = {}
    try:
        with open("/proc/self/status") as f:
            for line in f:
                if ":" not in line:
                    continue
                key, _, val = line.partition(":")
                key = key.strip()
                if key.startswith("Vm") or key.startswith("Rss"):
                    mem[key] = val.strip()
    except Exception as e:
        mem["error"] = str(e)
    return mem

def read_cgroup_memory():
    """
    cgroup v2 の memory.current（コンテナ全体の現在メモリ使用量 bytes）を取得。
    存在しない場合は None。memory.stat の anon / file / shmem / kernel も補足する。
    """
    result = {}
    for path, key in [
        (CGROUP_MEMORY_CURRENT, "memory_current_bytes"),
    ]:
        try:
            with open(path) as f:
                raw = f.read().strip()
            result[key] = int(raw)
        except FileNotFoundError:
            result[key] = None
            result[f"{key}_note"] = f"not found: {path}"
        except Exception as e:
            result[key] = None
            result[f"{key}_error"] = str(e)
    # memory.stat から anon / file を補足
    try:
        stat = {}
        with open(CGROUP_MEMORY_STAT) as f:
            for line in f:
                parts = line.strip().split()
                if len(parts) == 2 and parts[0] in {"anon", "file", "shmem", "kernel"}:
                    stat[parts[0]] = int(parts[1])
        result["memory_stat"] = stat
    except FileNotFoundError:
        result["memory_stat"] = None
    except Exception as e:
        result["memory_stat_error"] = str(e)
    return result

def vmrss_kb():
    """VmRSS の現在値 (kB) だけを高速に取得。"""
    try:
        with open("/proc/self/status") as f:
            for line in f:
                if line.startswith("VmRSS:"):
                    return int(line.split()[1])
    except Exception:
        pass
    return 0

@app.route("/ping")
def ping():
    return Response("OK", status=200)

@app.route("/invocations", methods=["POST"])
def invocations():
    try:
        raw = request.data or b"{}"
        body = json.loads(raw)
    except Exception:
        body = {}

    alloc_mb     = int(body.get("alloc_mb", 0))
    sleep_sec    = int(body.get("sleep_sec", 0))
    interval_sec = int(body.get("interval_sec", 10))

    # --- 確保前 ---
    before        = read_proc_status()
    cgroup_before  = read_cgroup_memory()

    # --- ヒープ確保（ページイン保証） ---
    data = None
    alloc_error = None
    alloc_actual_bytes = 0
    t_alloc_start = time.monotonic()
    if alloc_mb > 0:
        try:
            data = bytearray(os.urandom(alloc_mb * 1024 * 1024))
            alloc_actual_bytes = len(data)
            _ = data[0]
            _ = data[-1]
        except Exception as e:
            alloc_error = str(e)
    t_alloc_ms = round((time.monotonic() - t_alloc_start) * 1000, 1)

    # --- 確保後 ---
    after         = read_proc_status()
    cgroup_after  = read_cgroup_memory()

    # --- 解放 ---
    del data
    gc.collect()

    # --- 解放後 ---
    freed         = read_proc_status()
    cgroup_freed  = read_cgroup_memory()

    # --- 解放後スリープ ---
    sleep_snapshots = []
    if sleep_sec > 0:
        t_sleep_start = time.monotonic()
        next_snap = interval_sec
        while True:
            elapsed = time.monotonic() - t_sleep_start
            if elapsed >= sleep_sec:
                break
            if elapsed >= next_snap:
                snap_rss = vmrss_kb()
                snap_cgroup = read_cgroup_memory()
                sleep_snapshots.append({
                    "elapsed_sec": round(elapsed, 1),
                    "vmrss_kb": snap_rss,
                    "cgroup_memory_current_bytes": snap_cgroup.get("memory_current_bytes"),
                })
                next_snap += interval_sec
            time.sleep(0.5)
        # 最終スナップ
        sleep_snapshots.append({
            "elapsed_sec": round(time.monotonic() - t_sleep_start, 1),
            "vmrss_kb": vmrss_kb(),
            "cgroup_memory_current_bytes": read_cgroup_memory().get("memory_current_bytes"),
        })

    def kb_val(d, key):
        v = d.get(key, "0 kB")
        try:
            return int(v.split()[0])
        except Exception:
            return 0

    rss_before_kb = kb_val(before, "VmRSS")
    rss_after_kb  = kb_val(after,  "VmRSS")
    rss_freed_kb  = kb_val(freed,  "VmRSS")

    return jsonify({
        "alloc_mb_requested":  alloc_mb,
        "alloc_actual_bytes":  alloc_actual_bytes,
        "alloc_time_ms":       t_alloc_ms,
        "alloc_error":         alloc_error,
        "vmrss": {
            "before_kb":       rss_before_kb,
            "after_kb":        rss_after_kb,
            "freed_kb":        rss_freed_kb,
            "delta_alloc_kb":  rss_after_kb - rss_before_kb,
            "delta_freed_kb":  rss_freed_kb - rss_after_kb,
        },
        "proc_status": {
            "before": before,
            "after":  after,
            "freed":  freed,
        },
        # --- cgroup ---
        "cgroup_memory": {
            "before": cgroup_before,
            "after":  cgroup_after,
            "freed":  cgroup_freed,
        },
        # --- スリープ中の推移 ---
        "sleep_snapshots": sleep_snapshots,
    })

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=8080)