diff --git a/tests/test_modal_fn_volume.py b/tests/test_modal_fn_volume.py index 87b447b3..d523d87f 100644 --- a/tests/test_modal_fn_volume.py +++ b/tests/test_modal_fn_volume.py @@ -54,3 +54,35 @@ def test_hostmem_sampler(tmp_path): s = HostMemSampler(1, str(tmp_path / "h.jsonl")) rec = s.sample() assert rec["event"] == "modal_host_sample" and rec["meminfo_used"] > 0 and s.peak_host > 0 + + +def test_hostmem_cgroup_v2_and_v1_and_rss(tmp_path): + """s19-compaction-g3-20261010b: cgroup_current/peak were null on Modal (gVisor); + the sampler now finds the own cgroup (v2, else v1) and adds summed process RSS.""" + from yeto import modal_runner as mr + + v2 = tmp_path / "v2"; (v2 / "c").mkdir(parents=True) + (v2 / "c" / "memory.current").write_text("1000\n") + (v2 / "c" / "memory.peak").write_text("2000\n") + (v2 / "c" / "memory.max").write_text("max\n") + (v2 / "c" / "memory.events").write_text("low 0\nhigh 0\nmax 3\noom 1\noom_kill 1\n") + pc = tmp_path / "pc2"; pc.write_text("0::/c\n") + m = mr._cgroup_mem(str(v2), str(pc)) + assert m == {"cgroup_version": 2, "cgroup_current": 1000, "cgroup_peak": 2000, "cgroup_max": "max", + "cgroup_oom_kill": 1, "cgroup_oom": 1} + v1 = tmp_path / "v1"; (v1 / "memory" / "x").mkdir(parents=True) + d = v1 / "memory" / "x" + (d / "memory.usage_in_bytes").write_text("300\n"); (d / "memory.max_usage_in_bytes").write_text("400\n") + (d / "memory.limit_in_bytes").write_text("68719476736\n"); (d / "memory.failcnt").write_text("0\n") + (d / "memory.oom_control").write_text("oom_kill_disable 0\nunder_oom 0\noom_kill 2\n") + pc1 = tmp_path / "pc1"; pc1.write_text("4:memory:/x\n0::/\n") + m1 = mr._cgroup_mem(str(v1), str(pc1)) + assert m1["cgroup_version"] == 1 and m1["cgroup_current"] == 300 and m1["cgroup_max"] == 68719476736 + assert m1["cgroup_oom_kill"] == 2 + none = mr._cgroup_mem(str(tmp_path / "missing"), str(tmp_path / "nope")) + assert none["cgroup_version"] is None and none["cgroup_current"] is None + proc = tmp_path / "proc"; (proc / "12").mkdir(parents=True); (proc / "self").mkdir() + (proc / "12" / "status").write_text("Name:\tx\nVmRSS:\t 2048 kB\n") + assert mr._proc_rss_sum(str(proc)) == 2048 * 1024 + rec = mr._host_mem_used_bytes() # this machine: at least one memory source is non-null + assert rec["proc_rss_sum"] and (rec["cgroup_current"] is not None or rec["meminfo_used"]) diff --git a/yeto/modal_runner.py b/yeto/modal_runner.py index e9559e3c..d2becb1a 100644 --- a/yeto/modal_runner.py +++ b/yeto/modal_runner.py @@ -510,9 +510,109 @@ def container_command(run_script: str) -> list[str]: return ["bash", "-lc", f"cd {shlex.quote(CONTAINER_WORKDIR)} && {run_script}"] +CGROUP_ROOT = "/sys/fs/cgroup" + + +def _read_int(path: str): + try: + with open(path) as f: + raw = f.read().strip() + except Exception: # noqa: BLE001 + return None + if raw == "max": + return "max" + try: + return int(raw) + except ValueError: + return None + + +def _read_kv(path: str) -> dict: + out: dict = {} + try: + with open(path) as f: + for line in f: + parts = line.split() + if len(parts) == 2 and parts[1].lstrip("-").isdigit(): + out[parts[0]] = int(parts[1]) + except Exception: # noqa: BLE001 + pass + return out + + +def _own_cgroup_dirs(root: str, proc_cgroup: str) -> tuple[list[str], list[str]]: + """(v2 dirs, v1 memory dirs) to try, own cgroup first, then the root.""" + v2, v1 = [], [] + try: + with open(proc_cgroup) as f: + for line in f: + _, ctrl, rel = line.rstrip("\n").split(":", 2) + rel = rel.lstrip("/") + if ctrl == "": + v2.append(os.path.join(root, rel)) + elif "memory" in ctrl.split(","): + v1.append(os.path.join(root, "memory", rel)) + except Exception: # noqa: BLE001 + pass + v2.append(root) + v1.append(os.path.join(root, "memory")) + return v2, v1 + + +def _cgroup_mem(root: str = CGROUP_ROOT, proc_cgroup: str = "/proc/self/cgroup") -> dict: + """Container memory from cgroup v2 (memory.current/peak/max/events) or, + failing that, cgroup v1 (memory.usage_in_bytes/max_usage_in_bytes/ + limit_in_bytes, oom_control). Keys stay None when nothing is readable.""" + out = {"cgroup_version": None, "cgroup_current": None, "cgroup_peak": None, + "cgroup_max": None, "cgroup_oom_kill": None, "cgroup_oom": None} + v2, v1 = _own_cgroup_dirs(root, proc_cgroup) + for d in v2: + cur = _read_int(os.path.join(d, "memory.current")) + if isinstance(cur, int): + ev = _read_kv(os.path.join(d, "memory.events")) + out.update(cgroup_version=2, cgroup_current=cur, + cgroup_peak=_read_int(os.path.join(d, "memory.peak")), + cgroup_max=_read_int(os.path.join(d, "memory.max")), + cgroup_oom_kill=ev.get("oom_kill"), cgroup_oom=ev.get("oom")) + return out + for d in v1: + cur = _read_int(os.path.join(d, "memory.usage_in_bytes")) + if isinstance(cur, int): + ctl = _read_kv(os.path.join(d, "memory.oom_control")) + out.update(cgroup_version=1, cgroup_current=cur, + cgroup_peak=_read_int(os.path.join(d, "memory.max_usage_in_bytes")), + cgroup_max=_read_int(os.path.join(d, "memory.limit_in_bytes")), + cgroup_oom_kill=ctl.get("oom_kill"), cgroup_oom=_read_int(os.path.join(d, "memory.failcnt"))) + return out + return out + + +def _proc_rss_sum(proc: str = "/proc") -> int | None: + """Sum of VmRSS over every visible process (container view; cgroup-free fallback).""" + total, seen = 0, 0 + try: + names = os.listdir(proc) + except Exception: # noqa: BLE001 + return None + for name in names: + if not name.isdigit(): + continue + try: + with open(os.path.join(proc, name, "status")) as f: + for line in f: + if line.startswith("VmRSS:"): + total += int(line.split()[1]) * 1024 + seen += 1 + break + except Exception: # noqa: BLE001 + continue + return total if seen else None + + def _host_mem_used_bytes() -> dict: - """Container-visible host memory: /proc/meminfo (MemTotal - MemAvailable) - plus cgroup memory.current/peak when readable (None otherwise).""" + """Container-visible host memory: /proc/meminfo (MemTotal - MemAvailable), + the container's cgroup memory (v2, else v1; incl. limit and OOM-kill + count) and the summed RSS of all visible processes. Unreadable -> None.""" out: dict = {} try: info = {} @@ -524,13 +624,8 @@ def _host_mem_used_bytes() -> dict: out["meminfo_used"] = info.get("MemTotal", 0) - info.get("MemAvailable", 0) except Exception: # noqa: BLE001 - advisory pass - for key, path in (("cgroup_current", "/sys/fs/cgroup/memory.current"), - ("cgroup_peak", "/sys/fs/cgroup/memory.peak")): - try: - with open(path) as f: - out[key] = int(f.read().strip()) - except Exception: # noqa: BLE001 - out[key] = None + out.update(_cgroup_mem()) + out["proc_rss_sum"] = _proc_rss_sum() return out @@ -579,7 +674,8 @@ def sample(self) -> dict: except Exception: # noqa: BLE001 rec["gpu_mem_used_mib"] = None rec["gpu_util_pct"] = None - used = max(v or 0 for v in (rec.get("meminfo_used"), rec.get("cgroup_current"))) + used = max(v if isinstance(v, int) else 0 + for v in (rec.get("meminfo_used"), rec.get("cgroup_current"))) self.peak_host = max(self.peak_host, used) for i, v in enumerate(rec["gpu_mem_used_mib"] or []): if i >= len(self.peak_gpu):