Skip to content

feat: GPU profiling kit and capacity note template - #4

Open
matthewberry wants to merge 3 commits into
mainfrom
mjberry/feat/gpu-profiling-kit
Open

matthewberry wants to merge 3 commits into
mainfrom
mjberry/feat/gpu-profiling-kit

Conversation

@matthewberry

@matthewberry matthewberry commented Oct 20, 2026 •

Copy link
Copy Markdown
Contributor

Part of #2. Everything except the last criterion, the run on the mmli2 GPU node, which waits for a go-ahead (commands below).

What

profiling/, a kit for the profile that ends every tool's compute phase:

  • gpu_profile.py run: runs a series of input sizes in one Job in staging (ConfigMap plus Job from job.yaml, created rather than applied), follows the pod log so records survive a deadline kill, re-checks for GPU contention every 10 s, cleans up, and writes profile.json and profile.md. Per size it records wall time, card VRAM baseline and peak (nvidia-smi every 200 ms), peak memory from the cgroup (counting workers and /dev/shm), and the single largest process's peak RSS. It fits fixed + per_item × size for time and for VRAM over the baseline.
  • job.yaml: mirrors how the backend launches a GPU tool: the nvidia.com/gpu.present selector, the nvidia.com/gpu toleration, runtimeClassName: nvidia, a 1-GPU request, root, the /root/.cache mount on mmli-shared-job-data, and /dev/shm only when --shm is given (as the backend does for shmSizeLimit). Requests and limits are set separately, defaulting to ml-simplefold's. It runs the same python -m wrapper.entrypoint on a job.json. Each size is a fresh process, so model load counts as fixed cost every time. One unmeasured warm-up run comes first, so a cold weight cache doesn't skew the first point.
  • Guard (gpu_profile.py check, run before every run): it blocks on any pod that holds or waits for nvidia.com/gpu, and on any Job pod under the nvidia runtime without a GPU request (cleandb-mepesm and six others, which share the card unqueued). Resident non-Job pods on the nvidia runtime (the gpu-operator DaemonSets and chemscraper-services) don't block; they're part of the baseline.
  • inputs.py: make_job(size), the per-tool hook. Here it's implemented for the sample tool.
  • CAPACITY_NOTE.md plus gpu_profile.py capacity: mean GPU queue wait and utilization at 1×, 5× and 10× traffic, treating the GPU as one FIFO M/G/1 queue (Pollaczek–Khinchine) over the GPU-requesting tools.
  • profiling/README.md: how to run it, and how to derive maxItems (the smaller of the time limit, the VRAM limit, and the largest size that succeeded) and the job deadline from the profile.
  • CI runs the kit's tests (python -m unittest discover -s profiling; standard library only).

Found along the way (not changed here)

  • The backend can't set a deadline per tool. create_job reads activeDeadlineSeconds only from kubernetes_jobs.defaults (259200 s, 3 days). A hung GPU job can hold the only GPU for 3 days. The README says so, and the capacity note records the derived value until the backend supports one.
  • Most GPU tools never ask the scheduler for the GPU. From mmli-backend's chart/values.mmli2.*.yaml, only ml-simplefold (through the base values), ezspec-unidock and ezspec-inference request nvidia.com/gpu. Seven more (cleandb-mepesm and six others) run under the nvidia runtime without a request.

How it was checked

  • 23 unit tests pass locally: the blocker rules, the fit, the report, the manifest, the queue formula, and the in-pod driver run end to end against the sample tool (no GPU), including a failing input.
  • The driver ran inside the built sample image (python:3.12-slim) on sizes 1, 10 and 100.
  • gpu_profile.py render … | kubectl apply --dry-run=client --validate=strict -f - passes.
  • gpu_profile.py check against mmli2 (read-only): no blockers, with the six resident pods listed.
  • Not done: an actual run on the GPU node. Nothing was created on any cluster.

The GPU run (criterion 5), pending a go-ahead

cd profiling
python3 gpu_profile.py check
python3 gpu_profile.py run --image ghcr.io/moleculemaker/job-template:sha-3a5c0da --tool job-template --sizes 1,10,100 --out out-job-template

It creates and then deletes gpu-profile-job-template-<ts> (a ConfigMap and a Job) in staging, holds the GPU for about a minute, and creates uws/jobs/job-template/.cache on the shared volume. The sample tool is CPU-only, so this checks the mechanics (scheduling, the nvidia-smi sampling, logs, the report), not real GPU numbers. VRAM over the baseline should come out near 0.

profiling/gpu_profile.py runs a series of input sizes in one Job on the mmli2
GPU node, mirroring how the backend launches a GPU tool, and records wall
time, peak VRAM over the card's baseline and peak RSS per size, then fits
fixed vs per-item cost. It refuses to launch while any pod holds or waits for
the GPU, or a Job uses the nvidia runtime without requesting it.
CAPACITY_NOTE.md and 'gpu_profile.py capacity' project GPU queue wait at
1x/5x/10x traffic (M/G/1); the README derives maxItems and the job deadline
from the profile.

Part of #2

@matthewberry matthewberry left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review of the GPU profiling kit. Nothing was run. Each finding comes from reading the diff, plus a comparison with mmli-backend origin/main (job.tpl.yaml, chart/values*.yaml). Treat them as hypotheses.

The guard's "GPU busy" rule holds for what it checks. It skips pods that are neither Pending nor Running. It counts any pod that requests nvidia.com/gpu, and any Job-owned pod with runtimeClassName: nvidia, whether Pending or Running, as a blocker. Its gaps are about when it checks (finding 3) and what it can't see (finding 4).

  1. profiling/gpu_profile.py:142: one ConfigMap, applied client-side. Every size's job.json goes into one ConfigMap, applied with kubectl apply. That copies the whole object into the last-applied annotation, which is capped at 262144 bytes. Realistic inputs, such as sizes 1–500 of ~500-residue proteins, go past that, and the apply fails with metadata.annotations: Too long. So the profile can't reach the sizes it exists to find. Use kubectl create, server-side apply, or the shared volume.
  2. profiling/gpu_profile.py:145: a deadline hit loses every record silently. When activeDeadlineSeconds kills the Job, the pod is gone. kubectl logs then returns nothing, and nobody checks its exit code. parse_log returns [], and the report says fewer than two sizes succeeded, even though the small sizes finished. There's also no per-size timeout.
  3. profiling/gpu_profile.py:136: the guard runs once, before the apply. A runtime-class tool can start mid-run: cleandb-mepesm, aceretro, novostoic-enzrank, oed-*, reactionminer. It then adds VRAM and compute contention to later sizes, and the fit absorbs that as per-item cost. Re-running gpu_blockers while waiting or after the run, and marking the run tainted, would catch it.
  4. profiling/gpu_profile.py:58: resident services aren't constant. chemscraper-services' GPU use depends on chemscraper jobs that run on CPU worker nodes, which the guard never sees. A spike inside one size's window counts as that tool's VRAM and time, because the driver uses max(during) minus min(baseline).
  5. profiling/CAPACITY_NOTE.md:45: service_s is too short. It is defined as the tool process's run time. But a GPU-requesting pod holds the GPU for its whole life: image pull (imagePullPolicy: Always), the init download, and post-job. Underestimating rho makes the Pollaczek–Khinchine wait look bounded at 5× and 10× when it may not be. Use the pod's scheduled-to-terminated time instead.
  6. profiling/driver.py:90: ru_maxrss is the peak of one process, not the sum. It also leaves out /dev/shm, which counts against the cgroup limit. A tool with 4×3 GiB DataLoader workers plus a 6 GiB main process reports ~6 GiB while the cgroup peaks above 18 GiB, and a memory limit derived from that number gets OOMKilled. Read the cgroup's memory.peak instead.
  7. profiling/job.yaml:84: /dev/shm is always 8Gi. The backend only mounts dshm when values set shmSizeLimit (today only ezspec-inference, at 4Gi). Every other tool gets 64Mi. So the profile can pass a PyTorch tool that crashes in prod with "insufficient shared memory". Default to none, or mirror the tool's values.
  8. profiling/job.yaml:70: the cache mount doesn't match the existing tools. ml-simplefold mounts .cache at /root/.cache/torch/hub, and cleandb-mepesm and ezspec-inference mount .hub under torch/hub. Mounting at /root/.cache nests downloads inside prod's hub dir, and it masks weights baked into the image.
  9. profiling/job.yaml:49: it doesn't "mirror" the backend's resources or entrypoint. It defaults to requests==limits of 4 CPU / 16Gi. The backend gives ml-simplefold, for example, 4/24Gi requests and 16/32Gi limits, so the profile OOMKills sizes that run in prod. command: python /profile/driver.py also bypasses an image ENTRYPOINT that may activate a conda env.
  10. profiling/gpu_profile.py:54: two doc errors. "cleandb-mepesm and seven other tools" is wrong: values.mmli2.prod.yaml has six others, and the PR body says both six and eight. Also, CAPACITY_NOTE.md:29 says these tools' memory "goes in the budget above", but no budget appears above it.

- create, not apply: client-side apply copies the ConfigMap into a 256 KiB
  annotation (469 KB for sizes up to 500 of 500-residue inputs). Inputs over
  a ConfigMap's ~1 MiB are refused before anything is created.
- Follow the pod log from the start, so records survive a deadline kill;
  profile.md lists sizes that didn't finish.
- Re-check for GPU jobs every 10 s while running; profile.json records them
  under contended_by and profile.md marks the profile contended.
- Memory peak from the cgroup (mem_peak_mib), which counts workers and
  /dev/shm; ru_maxrss reported 213 MiB for three concurrent 200 MiB workers.
- /dev/shm only with --shm, as the backend mounts it only for shmSizeLimit.
- Separate requests and limits, defaulting to ml-simplefold's.
- Docs: service time is the whole pod's GPU hold; VRAM baseline caveat;
  six other runtime-only tools, not seven; the ENTRYPOINT assumption.
@matthewberry

Copy link
Copy Markdown
Contributor Author

Replies to the review. The fixes are all in bac0ba9. Tests: 23 → 32, all passing (python -m unittest discover -s profiling). The new run() tests drive it against a fake kubectl on PATH. The manifest passes kubectl apply --dry-run=client --validate=strict. Nothing was created on the cluster.

  1. One ConfigMap, applied client-side: reproduced, fixed. render for sizes 1–500 of 500-residue inputs, through kubectl apply --dry-run=client -o json, gives a last-applied annotation of 469,168 bytes, against the 262,144 limit. run now uses kubectl create, and refuses inputs over a ConfigMap's ~1 MiB before creating anything (test: test_inputs_too_large_for_a_configmap_are_refused_before_anything_is_created). Above 1 MiB, the README says to split the series across runs.
  2. A deadline hit loses every record: fixed. Not reproduced on the cluster (that needs a write); verified from the code path. run read the log only after the Job ended, and the pod is gone after a deadline kill. It now follows the log from the start (kubectl logs -f … --pod-running-timeout), and profile.md lists the sizes that didn't finish (test_records_streamed_before_a_deadline_kill_are_kept). I didn't add a separate per-size timeout: with streaming, the finished sizes survive, and the unfinished one is named.
  3. The guard runs once: fixed. _wait re-checks every 10 s, excluding our own app=gpu-profile pods. Anything it sees goes in profile.json under contended_by, and profile.md marks the run contended (test_a_gpu_job_that_starts_mid_run_taints_the_profile, test_its_own_pod_is_not_contention). Before launch, a leftover profiling pod still blocks.
  4. Resident services aren't constant: documented, no code change. The card-level sampler can't attribute a chemscraper-services spike to a process (PIDs inside the container don't match the host's). The README has a new "Limits of the measurements" section saying so, and says to re-run a size that's out of line with its neighbors.
  5. service_s is too short: agreed, fixed in CAPACITY_NOTE.md. It now defines service time as the pod's PodScheduled → last finishedAt. For the new tool, it's the profiled wall_s plus the pull/init/post-job overhead measured on an existing GPU tool's pods. It also says why leaving the overhead out understates ρ at 5× and 10×.
  6. ru_maxrss is one process: reproduced, fixed. Three concurrent 200 MiB workers under os.wait4 give ru_maxrss = 213 MiB (python:3.11 container). The driver now samples the cgroup's memory.current (falling back to v1's memory.usage_in_bytes) into mem_peak_mib, and profile.md shows both columns. The README says to size memory from the cgroup figure (test_peak_memory_comes_from_the_cgroup_so_workers_count). I sampled rather than reading memory.peak, because that's container-lifetime and can't be reset per size on every kernel.
  7. /dev/shm is always 8Gi: confirmed, fixed. shmSizeLimit is set only for ezspec-inference (4Gi, chart/values.yaml:324 on mmli-backend origin/main). The mount now appears only with --shm (test_no_shm_volume_unless_asked_like_the_backend).
  8. The cache mount: declined. The kit is for tools built from this template, and the template's README (lines 97–101) and Dockerfile comment prescribe exactly this: uws/jobs/<tool>/.cache mounted at /root/.cache. The existing tools' torch/hub mounts predate that convention, and the profile should mirror what the new tool will get in prod. That includes masking anything baked under /root/.cache, which prod would do too.
  9. Resources and entrypoint: fixed and documented. Requests and limits are now separate (--cpu/--memory and --cpu-limit/--memory-limit), defaulting to ml-simplefold's 4/16Gi and 16/32Gi (test_requests_and_limits_are_set_separately). The command matches the template's ENTRYPOINT ["python", "-m", "wrapper.entrypoint"]. job.yaml now says to change the command if a tool's image changes that ENTRYPOINT.
  10. Doc errors: fixed. On mmli-backend origin/main, values.mmli2.prod.yaml has 10 tools on the nvidia runtime. Three request the GPU (ezspec-inference, ezspec-unidock, ml-simplefold), and seven don't: aceretro, cleandb-mepesm, novostoic-enzrank, oed-catpred, oed-dlkcat, oed-unikp, reactionminer. That's "cleandb-mepesm and six others", now consistent in the docstring and the PR body. The capacity note's "budget above" now points to "From profile to limits" in profiling/README.md.

@matthewberry matthewberry left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Re-review of bac0ba9. I read the diff and the dispositions comment. I did not run the kit's tests in this pass.

The fixes for 1, 3, 5, 6, 7, 9 and 10 look right:

  • kubectl create doesn't write the last-applied annotation.
  • The size cap runs before anything is created.
  • The pod template carries app: gpu-profile, so _other_pods excludes the run's own pod correctly.
  • The cgroup sampler counts workers and /dev/shm.
  • /dev/shm is mounted only with --shm.
  • Requests and limits are now set separately.

I accept the declines: 4 is documented, and 8 follows the template's own convention.

New points:

  1. profiling/__pycache__/*.cpython-315.pyc: bytecode is committed. bac0ba9 adds four of these files (driver, gpu_profile, inputs, test_gpu_profile), and the repo has no .gitignore entry for __pycache__. Remove them and add the ignore.
  2. profiling/gpu_profile.py run(): the log follower can still lose records silently. Fix 2 was verified only by reading the code, and I've only read it too:
    • stderr goes to DEVNULL. If kubectl logs -f job/<name> fails to attach (e.g. it gives up before the pod exists, or the pod is never Running), nothing reports it. log comes back empty, and the report says no sizes finished, which is finding 2's symptom again. Capturing stderr and printing it when no records come back would make that visible.
    • logs.communicate(timeout=120) raises TimeoutExpired if the stream doesn't close within 120 s of the Job ending. The finally deletes the Job, but profile.json/profile.md are never written, and the kubectl process is never killed. Catching the timeout, killing the process and parsing what was read would keep the records.
    • Please confirm on staging that kubectl logs -f job/X --pod-running-timeout=… started while the pod is still Pending waits for it, with your kubectl version.
  3. Minor, _other_pods: it excludes every app=gpu-profile pod in every namespace. A second person's profiling run that starts mid-run isn't recorded as contention. Filtering on this run's own Job name (the job-name label) instead would keep it.

…e own-pod exclusion

- kubectl logs stderr is captured: a failed attach is reported as log_error
  in profile.json and profile.md instead of looking like a run with no records.
- If the stream outlives the Job by LOG_DRAIN_TIMEOUT_S, kill it and keep
  what arrived, rather than raising before the profile is written.
- Contention polling skips only this run's pods (job-name label), so another
  profiling run counts.
- Untrack __pycache__ committed in bac0ba9; add .gitignore.
@matthewberry

Copy link
Copy Markdown
Contributor Author

Replies to the second pass. All four are fixed in e7600c0. Tests: 32 → 35, all passing. I wrote the three new run() tests first, against the fake kubectl, and confirmed they failed. Nothing was created on the cluster.

  1. __pycache__ committed: confirmed, fixed. bac0ba9 added four profiling/__pycache__/*.cpython-315.pyc files. They're untracked now, and a new .gitignore has __pycache__/ (the repo had none). git ls-files | grep -c pycache → 0.
  2. The log follower's stderr went to DEVNULL: fixed. stderr is captured. A non-zero exit from kubectl logs is written as log_error in profile.json, and as a "Log incomplete" line at the top of profile.md (test_a_failed_log_attach_is_reported_not_swallowed).
  3. communicate(timeout=…) timing out: fixed. On TimeoutExpired, the stream is killed, communicate() is called again to collect what had arrived, and the profile is written with a timed-out log_error. The timeout is now LOG_DRAIN_TIMEOUT_S, so the test can shorten it (test_a_log_stream_that_outlives_the_job_is_killed_and_what_arrived_is_kept).
  4. _other_pods excluded every profiling pod: fixed. It now excludes only pods labeled job-name=<this run's Job>, the label the Job controller puts on its pods. Another profiling run counts as contention (test_another_profiling_run_is_contention). This also removes the now-unused PROFILE_LABEL constant.

@matthewberry matthewberry left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Round 3 (e7600c0): all four round-2 findings are fixed. I read the diff only; I didn't run the tests.

  • The four .pyc files are untracked, and .gitignore now covers __pycache__/.
  • kubectl's stderr is now captured. A non-zero exit is recorded as log_error in profile.json and shown in profile.md.
  • If the log stream hasn't closed LOG_DRAIN_TIMEOUT_S after the Job ends, it is killed, the records already read are kept, and the profile is still written. Skipping the returncode check when log_error is already set correctly keeps the kill's negative returncode from being reported twice.
  • _other_pods(name) now excludes only pods labeled job-name=<name>, which the Job controller sets, so another profiling run counts as contention.

No new findings. One note: when a deadline kill deletes the pod, kubectl logs -f may exit non-zero, and that run will now show "Log incomplete" next to its "Did not finish" sizes. That's accurate, so no change is needed.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant