diff --git a/.agents/skills/senpai-status-check/SKILL.md b/.agents/skills/senpai-status-check/SKILL.md index 21329b6ff..efd7de353 100644 --- a/.agents/skills/senpai-status-check/SKILL.md +++ b/.agents/skills/senpai-status-check/SKILL.md @@ -114,16 +114,39 @@ PY # Fleet overview and real training process sweep kubectl --context pai-2 get pods -l app=senpai -o wide -for podref in $(kubectl --context pai-2 get pods -l app=senpai,role=student -o name | sort); do - pod=${podref#pod/} - student=${pod#senpai-}; student=${student%-*}; student=${student%-*} - out=$(kubectl --context pai-2 exec "$pod" -- sh -lc "ps -eo pid,etime,comm,args | awk '\$3 ~ /python/ && /train.py/ {print; exit}'" 2>/dev/null || true) +kubectl --context pai-2 get pods -l app=senpai,role=student -o json > /tmp/student_pods.json +uv run python - <<'PY' /tmp/student_runtime.tsv +import json, sys +for item in json.load(sys.stdin).get("items", []): + meta = item.get("metadata", {}) + labels = meta.get("labels", {}) + annotations = meta.get("annotations", {}) + pod = meta.get("name", "") + names = [n.strip() for n in annotations.get("senpai/student-names", "").split(",") if n.strip()] + if not names and labels.get("student"): + names = [labels["student"]] + grouped = "student" not in labels + for student in names: + runner = f"/workspace/senpai-{student}" if grouped else "/workspace/senpai" + target = f"{runner}/target" + home = f"/workspace/home-{student}" if grouped else "/root" + print(f"{student}\t{pod}\t{target}\t{home}\t{runner}") +PY +while IFS="$(printf '\t')" read -r student pod target home runner; do + out=$(kubectl --context pai-2 exec "$pod" -- sh -lc ' + target="${1%/}" + runner="${2%/}" + . "$runner/k8s/senpai-processes.sh" + pid=$(senpai_training_pids_for_target "$target" target | head -1) + [ -n "$pid" ] || exit 0 + ps -p "$pid" -o pid=,etime=,comm=,args= 2>/dev/null + ' sh "$target" "$runner" 2>/dev/null || true) if [ -n "$out" ]; then printf "%-10s PYTRAIN %s\n" "$student" "$(printf "%s" "$out" | sed 's/^[[:space:]]*//; s/[[:space:]][[:space:]]*/ /g' | cut -c1-180)" else printf "%-10s NO_PYTRAIN\n" "$student" fi -done +done /dev/null | sort -nr | awk 'NR==1{sub(/^[^ ]+ /,\"\"); print}'); [ -n \"\$latest\" ] && tail -40 \"\$latest\"" sh "$home" +done /dev/null || true) +kubectl --context pai-2 get pods -l app=senpai,role=student -o json > /tmp/student_pods.json +uv run python - <<'PY' /tmp/student_runtime.tsv +import json, sys +for item in json.load(sys.stdin).get("items", []): + meta = item.get("metadata", {}) + labels = meta.get("labels", {}) + annotations = meta.get("annotations", {}) + pod = meta.get("name", "") + names = [n.strip() for n in annotations.get("senpai/student-names", "").split(",") if n.strip()] + if not names and labels.get("student"): + names = [labels["student"]] + grouped = "student" not in labels + for student in names: + runner = f"/workspace/senpai-{student}" if grouped else "/workspace/senpai" + target = f"{runner}/target" + home = f"/workspace/home-{student}" if grouped else "/root" + print(f"{student}\t{pod}\t{target}\t{home}\t{runner}") +PY +while IFS="$(printf '\t')" read -r student pod target home runner; do + out=$(kubectl --context pai-2 exec "$pod" -- sh -lc ' + target="${1%/}" + runner="${2%/}" + . "$runner/k8s/senpai-processes.sh" + pid=$(senpai_training_pids_for_target "$target" target | head -1) + [ -n "$pid" ] || exit 0 + ps -p "$pid" -o pid=,etime=,comm=,args= 2>/dev/null + ' sh "$target" "$runner" 2>/dev/null || true) if [ -n "$out" ]; then printf "%-10s PYTRAIN %s\n" "$student" "$(printf "%s" "$out" | sed 's/^[[:space:]]*//; s/[[:space:]][[:space:]]*/ /g' | cut -c1-180)" else printf "%-10s NO_PYTRAIN\n" "$student" fi -done +done /dev/null | sort -nr | awk 'NR==1{sub(/^[^ ]+ /,\"\"); print}'); [ -n \"\$latest\" ] && tail -40 \"\$latest\"" sh "$home" +done /dev/null 2>&1; then + printf '{}\n' > "$HOME/.claude.json" +fi echo "=== Claude config installed ===" ls "$HOME/.claude/skills/wandb-primary/SKILL.md" "$HOME/.claude/agents/researcher-agent.md" @@ -147,7 +150,7 @@ if [ -n "${EXTRA_INSTRUCTIONS_B64:-}" ]; then fi # Add "$KEY_INFO" (reminder of student names etc) to PROMPT -KEY_INFO=$'\n\n Key information:\n\n Students: '"$STUDENT_NAMES"' | GPUs per Student: '"$GPUS_PER_STUDENT"' | Tag: '"$RESEARCH_TAG"' | Target repo: '"$GH_REPO"' | Target base branch: '"${TARGET_REPO_BRANCH:-}"' | Advisor Branch: '"$ADVISOR_BRANCH"' | W&B entity/project: '"$WANDB_ENTITY"'/'"$WANDB_PROJECT"$'\n' +KEY_INFO=$'\n\n Key information:\n\n Students: '"$STUDENT_NAMES"' | GPU request: '"$GPUS_PER_STUDENT"' | Tag: '"$RESEARCH_TAG"' | Target repo: '"$GH_REPO"' | Target base branch: '"${TARGET_REPO_BRANCH:-}"' | Advisor Branch: '"$ADVISOR_BRANCH"' | W&B entity/project: '"$WANDB_ENTITY"'/'"$WANDB_PROJECT"$'\n' FULL_PROMPT="${PROMPT}"$'\n\n'"${KEY_INFO}" # Heartbeat prompt for polling diff --git a/k8s/entrypoint-student-group.sh b/k8s/entrypoint-student-group.sh new file mode 100755 index 000000000..0d8e4ae38 --- /dev/null +++ b/k8s/entrypoint-student-group.sh @@ -0,0 +1,83 @@ +#!/bin/bash + +# SPDX-FileCopyrightText: 2026 CoreWeave, Inc. +# SPDX-License-Identifier: Apache-2.0 +# SPDX-PackageName: senpai + +set -e +set -o pipefail + +echo "=== Senpai Student Group: ${STUDENT_GROUP_NAME:-group} ===" +echo "Students: ${STUDENT_NAMES}" + +repo_auth_url="$(printf '%s' "$REPO_URL" | sed "s#https://github.com/#https://${GITHUB_TOKEN}@github.com/#")" + +student_extra_instructions() { + local student="$1" + [ -n "${STUDENT_EXTRA_INSTRUCTIONS_B64_JSON_B64:-}" ] || return 0 + printf '%s' "$STUDENT_EXTRA_INSTRUCTIONS_B64_JSON_B64" | base64 -d | + STUDENT_FOR_EXTRA="$student" python3 -c ' +import json +import os +import sys + +mapping = json.load(sys.stdin) +print(mapping.get(os.environ["STUDENT_FOR_EXTRA"], "")) +' +} + +terminate_group() { + local pid + echo "=== Terminating grouped student processes ===" + for pid in "${student_pids[@]:-}"; do + [ -n "$pid" ] || continue + kill -TERM "$pid" 2>/dev/null || true + done + wait 2>/dev/null || true +} + +IFS=',' read -r -a raw_students <<< "$STUDENT_NAMES" +students=() +for raw_student in "${raw_students[@]}"; do + student="$(printf '%s' "$raw_student" | xargs)" + [ -n "$student" ] && students+=("$student") +done + +if [ "${#students[@]}" -eq 0 ]; then + echo "ERROR: STUDENT_NAMES is empty" >&2 + exit 2 +fi + +student_pids=() +trap terminate_group INT TERM + +for student in "${students[@]}"; do + ( + set -e + export STUDENT_NAME="$student" + export WORKDIR="/workspace/senpai-${student}" + export HOME="/workspace/home-${student}" + export SENPAI_GROUPED_STUDENT_POD=1 + export EXTRA_INSTRUCTIONS_B64 + + EXTRA_INSTRUCTIONS_B64="$(student_extra_instructions "$student")" + + mkdir -p "$HOME" + git clone --branch "$REPO_BRANCH" --single-branch --depth 1 --no-tags "$repo_auth_url" "$WORKDIR" + git -C "$WORKDIR" remote set-url origin "$REPO_URL" + cd "$WORKDIR" + bash k8s/entrypoint-student.sh + ) & + pid="$!" + student_pids+=("$pid") + echo "Started student ${student} as pid ${pid}" +done + +set +e +wait -n "${student_pids[@]}" +status=$? +set -e + +echo "=== A grouped student process exited with status ${status}; restarting the group pod ===" +terminate_group +exit "$status" diff --git a/k8s/entrypoint-student.sh b/k8s/entrypoint-student.sh index 2d4fdc13f..c90efa863 100644 --- a/k8s/entrypoint-student.sh +++ b/k8s/entrypoint-student.sh @@ -7,7 +7,7 @@ set -e set -o pipefail -WORKDIR="/workspace/senpai" +WORKDIR="${WORKDIR:-/workspace/senpai}" GH_HISTORY_SCOPE="${GH_HISTORY_SCOPE:-branch}" TARGET_REPO_BRANCH="${TARGET_REPO_BRANCH:-}" export SENPAI_ROLE="student" @@ -15,6 +15,10 @@ export TARGET_WORKDIR="$WORKDIR/$PROBLEM_DIR" GIT_CREDENTIAL_FILE="$WORKDIR/.git-credentials" SENPAI_PLUGIN="$WORKDIR/plugins/senpai" +if [ -x /usr/bin/nvidia-smi ]; then + export PATH="/usr/bin:$PATH" +fi + echo "=== Senpai Student: $STUDENT_NAME ===" echo "Runner repo: $REPO_URL (branch: $REPO_BRANCH)" echo "Target repo: $TARGET_REPO_URL (base branch: ${TARGET_REPO_BRANCH:-}; advisor branch: $ADVISOR_BRANCH)" @@ -62,6 +66,9 @@ fi # --- Install checked-in Claude Code config into user scope --- mkdir -p "$HOME/.claude" cp -a "$WORKDIR/.claude/." "$HOME/.claude/" +if ! python3 -m json.tool "$HOME/.claude.json" >/dev/null 2>&1; then + printf '{}\n' > "$HOME/.claude.json" +fi echo "=== Claude config installed ===" ls "$HOME/.claude/skills/wandb-primary/SKILL.md" "$HOME/.claude/agents/researcher-agent.md" @@ -91,7 +98,7 @@ if [ -n "${EXTRA_INSTRUCTIONS_B64:-}" ]; then PROMPT="${PROMPT}"$'\n\n# Finally, some additional instructions\n\n'"$(printf '%s' "$EXTRA_INSTRUCTIONS_B64" | base64 -d)" fi -KEY_INFO=$'\n\nKey information:\n\nStudent: '"$STUDENT_NAME"' | GPUs per Student: '"$GPUS_PER_STUDENT"' | Target repo: '"$GH_REPO"' | Target base branch: '"${TARGET_REPO_BRANCH:-}"' | Advisor Branch: '"$ADVISOR_BRANCH"' | W&B entity/project: '"$WANDB_ENTITY"'/'"$WANDB_PROJECT"$'\n' +KEY_INFO=$'\n\nKey information:\n\nStudent: '"$STUDENT_NAME"' | GPU request: '"$GPUS_PER_STUDENT"' | Target repo: '"$GH_REPO"' | Target base branch: '"${TARGET_REPO_BRANCH:-}"' | Advisor Branch: '"$ADVISOR_BRANCH"' | W&B entity/project: '"$WANDB_ENTITY"'/'"$WANDB_PROJECT"$'\n' FULL_PROMPT="${PROMPT}"$'\n\n'"${KEY_INFO}" HEARTBEAT_PROMPT="Continue your student loop using the assigned PRs and GitHub issues listed in the Student research state below. The entrypoint owns assignment polling; do not start persistent GitHub polling monitors. For active training, use sparse wakeups plus training_log_status; do not stream per-epoch logs into Monitor." diff --git a/k8s/install-weave-cc-plugin.sh b/k8s/install-weave-cc-plugin.sh index 857888d3c..31c764e0a 100755 --- a/k8s/install-weave-cc-plugin.sh +++ b/k8s/install-weave-cc-plugin.sh @@ -18,6 +18,11 @@ # already-token-authenticated HTTPS form set up in the deployment entrypoint. git config --global url."https://${GITHUB_TOKEN}@github.com/".insteadOf "git@github.com:" +mkdir -p "$HOME/.weave_claude_plugin" +if ! python3 -m json.tool "$HOME/.weave_claude_plugin/settings.json" >/dev/null 2>&1; then + rm -f "$HOME/.weave_claude_plugin/settings.json" +fi + WEAVE_PROJECT="${WANDB_ENTITY}/${WANDB_PROJECT}" \ weave-claude-plugin install --non-interactive diff --git a/k8s/launch.py b/k8s/launch.py index 51da90fa4..de6129430 100644 --- a/k8s/launch.py +++ b/k8s/launch.py @@ -7,6 +7,7 @@ """Launch senpai advisor and student agents as K8s resources.""" import base64 +import json import sys from dataclasses import dataclass from pathlib import Path @@ -49,7 +50,8 @@ class Args: names: str = "" # comma-separated student names (e.g. "frieren,fern") n_students: int = 4 # number of students to launch (ignored if --names is provided) student_prefix: str = "" # make assignment labels unique across parallel launches using the same base names - gpus_per_student: int = 8 # GPUs requested by each student pod + gpus_per_student: int = 8 # GPUs requested by each singleton student pod or packed student group + students_per_gpu_pod: int = 1 # logical student loops to pack into each GPU pod cpu_per_gpu: int = 15 # CPU requested per student GPU memory_gi_per_gpu: int = 120 # memory Gi requested per student GPU repo_url: str = "https://github.com/wandb/senpai.git" # git repo URL (senpai runner) @@ -116,6 +118,13 @@ def load_extra_instructions(extra_instructions: str) -> str: def build_extra_instructions(args: Args, tag: str, student_list: list[str]) -> str: students = ", ".join(student_list) target_base = args.target_repo_branch or "" + pod_packing = "" + if args.students_per_gpu_pod > 1: + pod_packing = ( + f"- This launch may pack up to `{args.students_per_gpu_pod}` logical students into one GPU pod. " + f"The pod requests `{args.gpus_per_student}` GPU(s) total, so coordinate heavy GPU work with " + "the advisor and other assigned students instead of assuming exclusive device access.\n" + ) isolation = f"""# Launch isolation and run-limit rules - This launch is scoped to research tag `{tag}`, advisor branch `{args.advisor_branch}`, and target base branch `{target_base}`. @@ -123,7 +132,7 @@ def build_extra_instructions(args: Args, tag: str, student_list: list[str]) -> s - Do not inspect, compare, summarize, cherry-pick, borrow from, or base decisions on any PR or branch outside `{args.advisor_branch}` and the assigned student PR branches for this launch. - Do not use unrelated experiment runs or historical results unless the human explicitly names them during this launch. - Students branch from `{args.advisor_branch}`. Do not rebase or retarget work onto unrelated branches. -- Treat `SENPAI_TIMEOUT_MINUTES` and `SENPAI_MAX_EPOCHS` as hard per-training-run bounds. Do not override them or continue a run past them. +{pod_packing}- Treat `SENPAI_TIMEOUT_MINUTES` and `SENPAI_MAX_EPOCHS` as hard per-training-run bounds. Do not override them or continue a run past them. """ user_extra = load_extra_instructions(args.extra_instructions) return isolation if not user_extra else isolation + "\n# Additional operator instructions\n\n" + user_extra @@ -133,51 +142,69 @@ def encoded_extra_instructions(args: Args, tag: str, student_list: list[str]) -> return base64.b64encode(build_extra_instructions(args, tag, student_list).encode()).decode() -def render_student(template: str, student_name: str, tag: str, secret_name: str, args: Args) -> str: - student_configmap_name = f"senpai-config-student-{tag}-{student_name}" - student_deployment_name = f"senpai-{tag}-{student_name}" +def student_chunks(student_list: list[str], chunk_size: int) -> list[list[str]]: + return [student_list[i:i + chunk_size] for i in range(0, len(student_list), chunk_size)] + + +def student_common_config_data(args: Args, tag: str) -> dict[str, str]: + return { + "REPO_URL": args.repo_url, + "REPO_BRANCH": args.repo_branch, + "TARGET_REPO_URL": args.target_repo_url, + "TARGET_REPO_BRANCH": args.target_repo_branch, + "GH_REPO": target_repo_slug(args.target_repo_url), + "RESEARCH_TAG": tag, + "GPUS_PER_STUDENT": str(args.gpus_per_student), + "WANDB_ENTITY": args.wandb_entity, + "WANDB_PROJECT": args.wandb_project, + "WANDB_MODE": "online", + "ADVISOR_BRANCH": args.advisor_branch, + "GH_HISTORY_SCOPE": args.gh_history_scope, + "SENPAI_ENABLE_HUMAN_ISSUES": "true" if args.human_issues else "false", + "SENPAI_TIMEOUT_MINUTES": str(args.timeout_minutes), + "SENPAI_MAX_EPOCHS": str(args.max_epochs), + "SENPAI_POLL_INTERVAL_S": str(args.poll_interval_s), + "SENPAI_POLL_JITTER_S": str(args.poll_jitter_s), + "STUDENT_CLAUDE_WATCHDOG_INTERVAL_S": str(args.student_claude_watchdog_interval_s), + "STUDENT_CLAUDE_WATCHDOG_JITTER_S": str(args.student_claude_watchdog_jitter_s), + "STUDENT_CLAUDE_MIN_RUNTIME_S": str(args.student_claude_min_runtime_s), + "STUDENT_CLAUDE_STALE_LOG_S": str(args.student_claude_stale_log_s), + "STUDENT_ASSIGNMENT_DRIFT_GRACE_S": str(args.student_assignment_drift_grace_s), + "PROBLEM_DIR": args.problem_dir, + "PVC_MOUNT_PATH": args.pvc_mount_path, + "SENPAI_START_GATE_PATH": args.start_gate_path, + } + + +def render_student_deployment( + template: str, + *, + deployment_name: str, + configmap_name: str, + id_label_key: str, + id_label_value: str, + student_names_csv: str, + container_name: str, + runner_workdir: str, + entrypoint: str, + tag: str, + secret_name: str, + args: Args, +) -> str: student_cpu = args.cpu_per_gpu * args.gpus_per_student student_memory_gi = args.memory_gi_per_gpu * args.gpus_per_student - configmap = render_configmap( - name=student_configmap_name, - labels={"app": "senpai", "role": "student", "research-tag": tag}, - data={ - "REPO_URL": args.repo_url, - "REPO_BRANCH": args.repo_branch, - "TARGET_REPO_URL": args.target_repo_url, - "TARGET_REPO_BRANCH": args.target_repo_branch, - "GH_REPO": target_repo_slug(args.target_repo_url), - "STUDENT_NAME": student_name, - "RESEARCH_TAG": tag, - "GPUS_PER_STUDENT": str(args.gpus_per_student), - "WANDB_ENTITY": args.wandb_entity, - "WANDB_PROJECT": args.wandb_project, - "WANDB_MODE": "online", - "ADVISOR_BRANCH": args.advisor_branch, - "GH_HISTORY_SCOPE": args.gh_history_scope, - "SENPAI_ENABLE_HUMAN_ISSUES": "true" if args.human_issues else "false", - "SENPAI_TIMEOUT_MINUTES": str(args.timeout_minutes), - "SENPAI_MAX_EPOCHS": str(args.max_epochs), - "SENPAI_POLL_INTERVAL_S": str(args.poll_interval_s), - "SENPAI_POLL_JITTER_S": str(args.poll_jitter_s), - "STUDENT_CLAUDE_WATCHDOG_INTERVAL_S": str(args.student_claude_watchdog_interval_s), - "STUDENT_CLAUDE_WATCHDOG_JITTER_S": str(args.student_claude_watchdog_jitter_s), - "STUDENT_CLAUDE_MIN_RUNTIME_S": str(args.student_claude_min_runtime_s), - "STUDENT_CLAUDE_STALE_LOG_S": str(args.student_claude_stale_log_s), - "STUDENT_ASSIGNMENT_DRIFT_GRACE_S": str(args.student_assignment_drift_grace_s), - "EXTRA_INSTRUCTIONS_B64": encoded_extra_instructions(args, tag, [student_name]), - "PROBLEM_DIR": args.problem_dir, - "PVC_MOUNT_PATH": args.pvc_mount_path, - "SENPAI_START_GATE_PATH": args.start_gate_path, - }, - ) - deployment = render_template(template, { - "STUDENT_DEPLOYMENT_NAME": student_deployment_name, - "STUDENT_CONFIGMAP_NAME": student_configmap_name, - "STUDENT_NAME": student_name, + return render_template(template, { + "STUDENT_DEPLOYMENT_NAME": deployment_name, + "STUDENT_CONFIGMAP_NAME": configmap_name, + "STUDENT_METADATA_ID_LABEL": f" {id_label_key}: {id_label_value}", + "STUDENT_SELECTOR_ID_LABEL": f" {id_label_key}: {id_label_value}", + "STUDENT_TEMPLATE_ID_LABEL": f" {id_label_key}: {id_label_value}", + "STUDENT_NAMES_CSV": student_names_csv, + "STUDENT_CONTAINER_NAME": container_name, + "STUDENT_RUNNER_WORKDIR": runner_workdir, + "STUDENT_ENTRYPOINT": entrypoint, "RESEARCH_TAG": tag, "IMAGE": args.image, - "ADVISOR_BRANCH": args.advisor_branch, "PVC_CLAIM_NAME": args.pvc_claim_name, "PVC_MOUNT_PATH": args.pvc_mount_path, "LAUNCH_SECRET_NAME": secret_name, @@ -185,6 +212,73 @@ def render_student(template: str, student_name: str, tag: str, secret_name: str, "STUDENT_MEMORY": f"{student_memory_gi}Gi", "GPUS_PER_STUDENT": str(args.gpus_per_student), }) + + +def render_student(template: str, student_name: str, tag: str, secret_name: str, args: Args) -> str: + student_configmap_name = f"senpai-config-student-{tag}-{student_name}" + student_deployment_name = f"senpai-{tag}-{student_name}" + data = student_common_config_data(args, tag) + data.update({ + "STUDENT_NAME": student_name, + "EXTRA_INSTRUCTIONS_B64": encoded_extra_instructions(args, tag, [student_name]), + }) + configmap = render_configmap( + name=student_configmap_name, + labels={"app": "senpai", "role": "student", "research-tag": tag}, + data=data, + ) + deployment = render_student_deployment( + template, + deployment_name=student_deployment_name, + configmap_name=student_configmap_name, + id_label_key="student", + id_label_value=student_name, + student_names_csv=student_name, + container_name="student", + runner_workdir="/workspace/senpai", + entrypoint="k8s/entrypoint-student.sh", + tag=tag, + secret_name=secret_name, + args=args, + ) + return configmap + "\n---\n" + deployment + + +def render_student_group(template: str, student_names: list[str], group_index: int, tag: str, secret_name: str, args: Args) -> str: + group_name = f"group-{group_index + 1}" + student_names_csv = ",".join(student_names) + configmap_name = f"senpai-config-student-{tag}-{group_name}" + deployment_name = f"senpai-{tag}-{group_name}" + extra_by_student = { + student: encoded_extra_instructions(args, tag, [student]) + for student in student_names + } + data = student_common_config_data(args, tag) + data.update({ + "STUDENT_GROUP_NAME": group_name, + "STUDENT_NAMES": student_names_csv, + "STUDENTS_PER_GPU_POD": str(args.students_per_gpu_pod), + "STUDENT_EXTRA_INSTRUCTIONS_B64_JSON_B64": base64.b64encode(json.dumps(extra_by_student).encode()).decode(), + }) + configmap = render_configmap( + name=configmap_name, + labels={"app": "senpai", "role": "student", "research-tag": tag}, + data=data, + ) + deployment = render_student_deployment( + template, + deployment_name=deployment_name, + configmap_name=configmap_name, + id_label_key="student-group", + id_label_value=group_name, + student_names_csv=student_names_csv, + container_name="student-group", + runner_workdir="/workspace/senpai-group", + entrypoint="k8s/entrypoint-student-group.sh", + tag=tag, + secret_name=secret_name, + args=args, + ) return configmap + "\n---\n" + deployment @@ -238,6 +332,8 @@ def main(): args = sp.parse(Args, config_path=str(SENPAI_CONFIG)) if min(args.gpus_per_student, args.cpu_per_gpu, args.memory_gi_per_gpu) < 1: sys.exit("ERROR: --gpus_per_student, --cpu_per_gpu, and --memory_gi_per_gpu must all be at least 1") + if args.students_per_gpu_pod < 1: + sys.exit("ERROR: --students_per_gpu_pod must be at least 1") validate_timing_args(args) if args.gh_history_scope not in {"branch", "repo", "fresh"}: sys.exit("ERROR: --gh_history_scope must be one of: branch, repo, fresh") @@ -301,14 +397,25 @@ def main(): ) # --- Deploy students --- - for name in student_list: - manifest = render_student(student_template, name, args.tag, secret_name, args) - if args.dry_run: - print(f"--- Student: {name} ---") - print(manifest) - print() - else: - kubectl_apply(manifest, f"student {name}") + if args.students_per_gpu_pod == 1: + for name in student_list: + manifest = render_student(student_template, name, args.tag, secret_name, args) + if args.dry_run: + print(f"--- Student: {name} ---") + print(manifest) + print() + else: + kubectl_apply(manifest, f"student {name}") + else: + for group_index, names in enumerate(student_chunks(student_list, args.students_per_gpu_pod)): + manifest = render_student_group(student_template, names, group_index, args.tag, secret_name, args) + group_name = f"group-{group_index + 1}" + if args.dry_run: + print(f"--- Student group: {group_name} ({', '.join(names)}) ---") + print(manifest) + print() + else: + kubectl_apply(manifest, f"student group {group_name}") advisor_student_list = student_list if args.advisor and not args.dry_run: @@ -326,14 +433,19 @@ def main(): if not args.dry_run: print(f"\nLaunched {len(student_list)} students: {', '.join(student_list)}") + if args.students_per_gpu_pod > 1: + group_count = len(student_chunks(student_list, args.students_per_gpu_pod)) + print(f"Packed into {group_count} GPU pod(s), up to {args.students_per_gpu_pod} students per pod") if args.advisor: print("Launched advisor pod") print(f"\nMonitor:") print(f" kubectl get deployments -l research-tag={args.tag}") if args.advisor: print(f" kubectl get deployment senpai-advisor-{args.tag}") - if student_list: + if student_list and args.students_per_gpu_pod == 1: print(f" kubectl logs -f deployment/senpai-{args.tag}-{student_list[0]}") + elif student_list: + print(f" kubectl logs -f deployment/senpai-{args.tag}-group-1") print(f"\nStop:") print(f" kubectl delete deployments,configmaps,secrets -l research-tag={args.tag}") diff --git a/k8s/launch_helpers.py b/k8s/launch_helpers.py index 4bd321767..4595e3470 100644 --- a/k8s/launch_helpers.py +++ b/k8s/launch_helpers.py @@ -78,11 +78,22 @@ def existing_student_names(tag: str) -> list[str]: "-l", f"app=senpai,role=student,research-tag={tag}", "-o", - 'jsonpath={range .items[*]}{.metadata.labels.student}{"\\n"}{end}', + "json", ], capture_output=True, text=True, check=True, ) - return [line for line in result.stdout.splitlines() if line] + names: list[str] = [] + for item in json.loads(result.stdout or "{}").get("items", []): + metadata = item.get("metadata", {}) + labels = metadata.get("labels", {}) + annotations = metadata.get("annotations", {}) + label_student = labels.get("student", "") + if label_student: + names.append(label_student) + continue + annotated_students = annotations.get("senpai/student-names", "") + names.extend(name.strip() for name in annotated_students.split(",") if name.strip()) + return list(dict.fromkeys(names)) def render_template(template: str, replacements: dict[str, str]) -> str: diff --git a/k8s/senpai-processes.sh b/k8s/senpai-processes.sh new file mode 100644 index 000000000..9c33cc387 --- /dev/null +++ b/k8s/senpai-processes.sh @@ -0,0 +1,103 @@ +#!/usr/bin/env sh + +# SPDX-FileCopyrightText: 2026 CoreWeave, Inc. +# SPDX-License-Identifier: Apache-2.0 +# SPDX-PackageName: senpai + +# Shared process ownership helpers for Senpai pods. A singleton student pod can +# use global process ownership; grouped pods should scope ownership to the +# logical student's target checkout. + +senpai_realpath_dir() { + (cd "$1" && pwd -P) 2>/dev/null || printf '%s\n' "${1%/}" +} + +senpai_pid_cwd_under() { + senpai_pid="$1" + senpai_target="$(senpai_realpath_dir "$2")" + senpai_cwd="$(readlink -f "/proc/$senpai_pid/cwd" 2>/dev/null || true)" + case "$senpai_cwd" in + "$senpai_target"|"$senpai_target"/*) return 0 ;; + *) return 1 ;; + esac +} + +senpai_pid_descends_from() { + senpai_pid="$1" + senpai_ancestor="$2" + [ -n "$senpai_pid" ] && [ -n "$senpai_ancestor" ] || return 1 + while [ -n "$senpai_pid" ] && [ "$senpai_pid" != "0" ] && [ "$senpai_pid" != "1" ]; do + [ "$senpai_pid" = "$senpai_ancestor" ] && return 0 + senpai_pid="$(ps -o ppid= -p "$senpai_pid" 2>/dev/null | tr -d ' ')" + done + return 1 +} + +senpai_pid_belongs_to_target() { + senpai_pid="$1" + senpai_target="$2" + senpai_scope="${3:-target}" + senpai_ancestor="${4:-}" + [ "$senpai_scope" = "global" ] && return 0 + senpai_pid_cwd_under "$senpai_pid" "$senpai_target" && return 0 + senpai_pid_descends_from "$senpai_pid" "$senpai_ancestor" +} + +senpai_process_pids_for_target() { + senpai_kind="$1" + senpai_target="$(senpai_realpath_dir "$2")" + senpai_scope="${3:-target}" + senpai_ancestor="${4:-}" + ps -eo pid=,comm=,args= | + while read -r senpai_candidate_pid senpai_comm senpai_args; do + case "$senpai_kind:$senpai_comm" in + train:python|train:python[0-9.]*|train:torchrun|train:pt_elastic) ;; + claude:claude) ;; + claude:*) case "$senpai_args" in *" claude "*|*/claude\ *|*/claude) ;; *) continue ;; esac ;; + *) continue ;; + esac + if [ "$senpai_kind" = "train" ]; then + case "$senpai_args" in *train.py*) ;; *) continue ;; esac + fi + senpai_pid_belongs_to_target "$senpai_candidate_pid" "$senpai_target" "$senpai_scope" "$senpai_ancestor" && + printf '%s\n' "$senpai_candidate_pid" + done +} + +senpai_training_pids_for_target() { + senpai_process_pids_for_target train "$@" +} + +senpai_claude_pids_for_target() { + senpai_process_pids_for_target claude "$@" +} + +senpai_gpu_pids_for_target() { + senpai_target="$(senpai_realpath_dir "$1")" + senpai_scope="${2:-target}" + if [ "$senpai_scope" = "global" ]; then + nvidia-smi --query-compute-apps=pid --format=csv,noheader,nounits 2>/dev/null | + awk 'NF {print}' + return 0 + fi + nvidia-smi --query-compute-apps=pid --format=csv,noheader,nounits 2>/dev/null | + while read -r senpai_candidate_pid; do + [ -n "$senpai_candidate_pid" ] || continue + senpai_pid_cwd_under "$senpai_candidate_pid" "$senpai_target" && printf '%s\n' "$senpai_candidate_pid" + done +} + +senpai_count_lines() { + awk 'NF {n++} END {print n+0}' +} + +senpai_student_activity_snapshot() { + senpai_target="$(senpai_realpath_dir "$1")" + senpai_scope="${2:-target}" + senpai_branch="$(git -C "$senpai_target" branch --show-current 2>/dev/null || true)" + senpai_pytrain="$(senpai_training_pids_for_target "$senpai_target" "$senpai_scope" | senpai_count_lines)" + senpai_gpu="$(senpai_gpu_pids_for_target "$senpai_target" "$senpai_scope" | senpai_count_lines)" + senpai_claude="$(senpai_claude_pids_for_target "$senpai_target" "$senpai_scope" | senpai_count_lines)" + senpai_dirty="$(git -C "$senpai_target" status --porcelain 2>/dev/null | wc -l | tr -d ' ')" + printf '%s\t%s\t%s\t%s\t%s\n' "$senpai_branch" "$senpai_pytrain" "$senpai_gpu" "$senpai_claude" "${senpai_dirty:-0}" +} diff --git a/k8s/student-claude-watchdog.sh b/k8s/student-claude-watchdog.sh index f0c1250a2..e58283f2d 100755 --- a/k8s/student-claude-watchdog.sh +++ b/k8s/student-claude-watchdog.sh @@ -17,6 +17,24 @@ STUDENT_CLAUDE_STALE_LOG_S="${STUDENT_CLAUDE_STALE_LOG_S:-1200}" STUDENT_CLAUDE_KILL_GRACE_S="${STUDENT_CLAUDE_KILL_GRACE_S:-15}" STUDENT_ASSIGNMENT_DRIFT_GRACE_S="${STUDENT_ASSIGNMENT_DRIFT_GRACE_S:-1800}" +. "${SENPAI_PROCESS_HELPERS:-${WORKDIR:-/workspace/senpai}/k8s/senpai-processes.sh}" + +student_grouped_pod() { + [ "${SENPAI_GROUPED_STUDENT_POD:-}" = "1" ] +} + +student_target_workdir() { + (cd "$TARGET_WORKDIR" && pwd -P) 2>/dev/null || printf '%s\n' "${TARGET_WORKDIR%/}" +} + +student_process_scope() { + if student_grouped_pod; then + printf 'target\n' + else + printf 'global\n' + fi +} + student_file_mtime_s() { stat -c %Y "$1" 2>/dev/null || stat -f %m "$1" } @@ -33,19 +51,20 @@ print(",".join(f"#{number}" for number in numbers)) } student_has_active_training() { - ps -eo pid,ppid,comm,args | - awk '$3 ~ /^(python[0-9.]*|torchrun)$/ && $0 ~ /train[.]py/ { found = 1 } END { exit !found }' + local claude_pid="${1:-}" + [ -n "$(student_training_pids "$claude_pid" | head -1)" ] } student_training_pids() { - ps -eo pid,ppid,comm,args | - awk '$3 ~ /^(python[0-9.]*|torchrun|pt_elastic)$/ && $0 ~ /train[.]py/ { print $1 }' + local claude_pid="${1:-}" + senpai_training_pids_for_target "$(student_target_workdir)" "$(student_process_scope)" "$claude_pid" } student_training_snapshot() { - ps -eo pid,ppid,etimes,pcpu,pmem,comm,args | - awk '$6 ~ /^(python[0-9.]*|torchrun|pt_elastic)$/ && $0 ~ /train[.]py/ { print }' | - head -20 + local claude_pid="${1:-}" pids + pids="$(student_training_pids "$claude_pid" | head -20 | tr '\n' ',')" + [ -n "$pids" ] || return 0 + ps -p "${pids%,}" -o pid,ppid,etimes,pcpu,pmem,comm,args 2>/dev/null || true } student_log_watchdog_trigger() { @@ -73,8 +92,8 @@ student_stop_claude_tree() { } student_stop_training_trees() { - local pids pid - pids=$(student_training_pids | sort -nr) + local claude_pid="${1:-}" pids pid + pids=$(student_training_pids "$claude_pid" | sort -nr) [ -n "$pids" ] || return 0 echo "=== Claude watchdog: stopping active train.py processes ===" @@ -92,10 +111,10 @@ student_stop_training_trees() { } student_comment_watchdog_stop() { - local reason="$1" start_numbers="$2" current_numbers="$3" snapshot body num + local reason="$1" start_numbers="$2" current_numbers="$3" claude_pid="${4:-}" snapshot body num [ -n "$start_numbers" ] || return 0 - snapshot=$(student_training_snapshot || true) + snapshot=$(student_training_snapshot "$claude_pid" || true) [ -n "$snapshot" ] || snapshot="(no active train.py process found at comment time)" body=$(cat </dev/null 2>&1 || return 1 pod=$( @@ -596,22 +602,83 @@ student_pr_looks_live() { -o jsonpath='{range .items[?(@.status.phase=="Running")]}{.metadata.name}{"\n"}{end}' 2>/dev/null | head -1 ) || return 1 - [ -n "$pod" ] || return 1 + if [ -n "$pod" ]; then + printf '%s\t%s\t%s\n' "$pod" "$(student_target_path_for_workdir /workspace/senpai)" "/workspace/senpai" + return 0 + fi + + kubectl get pods -l "app=senpai,role=student,research-tag=${tag}" -o json 2>/dev/null | + STUDENT_TO_FIND="$student" PROBLEM_DIR_FOR_STUDENT="${PROBLEM_DIR:-target}" python3 -c ' +import json +import os +import sys + +student = os.environ["STUDENT_TO_FIND"] +problem_dir = os.environ["PROBLEM_DIR_FOR_STUDENT"].strip("/") +data = json.load(sys.stdin) +for item in data.get("items", []): + if item.get("status", {}).get("phase") != "Running": + continue + metadata = item.get("metadata", {}) + annotations = metadata.get("annotations", {}) + names = [ + name.strip() + for name in annotations.get("senpai/student-names", "").split(",") + if name.strip() + ] + if student in names: + workdir = f"/workspace/senpai-{student}" + print("{}\t{}/{}\t{}".format(metadata.get("name", ""), workdir, problem_dir, workdir)) + sys.exit(0) +sys.exit(1) + ' +} + +student_runner_is_grouped() { + case "$1" in + /workspace/senpai-*) return 0 ;; + *) return 1 ;; + esac +} +student_pod_activity_snapshot() { + local pod="$1" target_path="$2" runner_workdir="$3" grouped=0 + student_runner_is_grouped "$runner_workdir" && grouped=1 kubectl exec "$pod" -- sh -lc ' - branch=$(git -C /workspace/senpai/target branch --show-current 2>/dev/null || true) - [ "$branch" = "$1" ] || exit 1 - - # A stale GitHub timestamp is not actionable if the pod is still doing - # useful work on this exact PR branch. Count either active training, - # active GPU use, or Claude Code editing an uncommitted checkout. - pytrain=$(ps -eo comm=,args= | awk '\''$1 ~ /^python/ && /train[.]py/ {n++} END{print n+0}'\'') - gpu=$(nvidia-smi --query-compute-apps=pid --format=csv,noheader,nounits 2>/dev/null | awk '\''NF{n++} END{print n+0}'\'') - claude=$(ps -eo comm=,args= | awk '\''$1 ~ /^claude$/ || /[ /]claude( |$)/ {n++} END{print n+0}'\'') - dirty=$(git -C /workspace/senpai/target status --porcelain 2>/dev/null | wc -l | tr -d " ") - - [ "$pytrain" -gt 0 ] || [ "$gpu" -gt 0 ] || { [ "$claude" -gt 0 ] && [ "${dirty:-0}" -gt 0 ]; } - ' sh "$head_ref" >/dev/null 2>&1 + target_path="${1%/}" + grouped="$2" + runner_workdir="${3%/}" + scope="global" + [ "$grouped" = "1" ] && scope="target" + . "$runner_workdir/k8s/senpai-processes.sh" + senpai_student_activity_snapshot "$target_path" "$scope" + ' sh "$target_path" "$grouped" "$runner_workdir" 2>/dev/null +} + +student_pr_looks_live() { + local student="$1" head_ref="$2" tag="${RESEARCH_TAG:-}" + local pod_info pod target_path runner_workdir snapshot branch pytrain gpu claude dirty + [ -n "$student" ] && [ -n "$head_ref" ] && [ -n "$tag" ] || return 1 + command -v kubectl >/dev/null 2>&1 || return 1 + + pod_info=$(student_pod_and_target_path "$student") || return 1 + pod=${pod_info%%$'\t'*} + pod_info=${pod_info#*$'\t'} + target_path=${pod_info%%$'\t'*} + runner_workdir=${pod_info#*$'\t'} + [ -n "$pod" ] && [ -n "$target_path" ] && [ -n "$runner_workdir" ] || return 1 + snapshot=$(student_pod_activity_snapshot "$pod" "$target_path" "$runner_workdir") || return 1 + branch=${snapshot%%$'\t'*} + snapshot=${snapshot#*$'\t'} + pytrain=${snapshot%%$'\t'*} + snapshot=${snapshot#*$'\t'} + gpu=${snapshot%%$'\t'*} + snapshot=${snapshot#*$'\t'} + claude=${snapshot%%$'\t'*} + dirty=${snapshot#*$'\t'} + + [ "$branch" = "$head_ref" ] || return 1 + [ "${pytrain:-0}" -gt 0 ] || [ "${gpu:-0}" -gt 0 ] || { [ "${claude:-0}" -gt 0 ] && [ "${dirty:-0}" -gt 0 ]; } } suppress_live_stale_wips() { @@ -660,7 +727,7 @@ print(json.dumps(filtered)) # list_student_pod_anomalies list_student_pod_anomalies() { local students_csv="$1" branch="$2" tag="${RESEARCH_TAG:-}" - local all_prs open_heads anomalies="" student pod snapshot current_branch pytrain gpu expected_heads + local all_prs open_heads anomalies="" student pod snapshot current_branch pytrain gpu expected_heads pod_info target_path runner_workdir claude dirty # This helper runs inside advisor pods. Local/dev shells without Kubernetes # context should stay quiet rather than making every advisor poll noisy. [ -n "$tag" ] && command -v kubectl >/dev/null 2>&1 || { echo "[]"; return 0; } @@ -685,26 +752,22 @@ for pr in json.load(sys.stdin): student="${student//[[:space:]]/}" [ -n "$student" ] || continue - pod=$( - kubectl get pods -l "app=senpai,research-tag=${tag},student=${student}" \ - -o jsonpath='{range .items[?(@.status.phase=="Running")]}{.metadata.name}{"\n"}{end}' 2>/dev/null | - head -1 - ) || continue - [ -n "$pod" ] || continue - - snapshot=$(kubectl exec "$pod" -- sh -lc ' - branch=$(git -C /workspace/senpai/target branch --show-current 2>/dev/null || true) - # Only active training/GPU use is a zombie risk. Claude editing a - # checkout is handled by stale-WIP suppression above, not here. - pytrain=$(ps -eo comm=,args= | awk '\''$1 ~ /^python/ && /train[.]py/ {n++} END{print n+0}'\'') - gpu=$(nvidia-smi --query-compute-apps=pid --format=csv,noheader,nounits 2>/dev/null | awk '\''NF{n++} END{print n+0}'\'') - printf "%s\t%s\t%s\n" "$branch" "$pytrain" "$gpu" - ' 2>/dev/null) || continue + pod_info=$(student_pod_and_target_path "$student") || continue + pod=${pod_info%%$'\t'*} + pod_info=${pod_info#*$'\t'} + target_path=${pod_info%%$'\t'*} + runner_workdir=${pod_info#*$'\t'} + [ -n "$pod" ] && [ -n "$target_path" ] && [ -n "$runner_workdir" ] || continue + snapshot=$(student_pod_activity_snapshot "$pod" "$target_path" "$runner_workdir") || continue current_branch=${snapshot%%$'\t'*} snapshot=${snapshot#*$'\t'} pytrain=${snapshot%%$'\t'*} - gpu=${snapshot#*$'\t'} + snapshot=${snapshot#*$'\t'} + gpu=${snapshot%%$'\t'*} + snapshot=${snapshot#*$'\t'} + claude=${snapshot%%$'\t'*} + dirty=${snapshot#*$'\t'} if [ "${pytrain:-0}" -eq 0 ] && [ "${gpu:-0}" -eq 0 ]; then continue fi diff --git a/scripts/arm_senpai_cluster_cutoff.sh b/scripts/arm_senpai_cluster_cutoff.sh index e0e7f5ad4..8331f62e2 100755 --- a/scripts/arm_senpai_cluster_cutoff.sh +++ b/scripts/arm_senpai_cluster_cutoff.sh @@ -3,9 +3,9 @@ # # The job waits inside Kubernetes until all expected Senpai pods are Ready, # records that timestamp on the PVC, sleeps for the requested budget, harvests -# Claude Code conversation logs from /root/.claude, writes them to the PVC, then -# deletes the tagged Senpai deployments/configmaps/secrets. This keeps the hard -# cutoff independent of the operator laptop staying awake. +# Claude Code conversation logs from singleton and grouped student homes, writes +# them to the PVC, then deletes the tagged Senpai deployments/configmaps/secrets. +# This keeps the hard cutoff independent of the operator laptop staying awake. set -euo pipefail @@ -250,17 +250,23 @@ sleep_until() { } copy_one_pod() { - local pod="$1" tag="$2" dest="$3" pod_dir + local pod="$1" tag="$2" student_names="$3" dest="$4" pod_dir pod_dir="${dest}/pods/$(safe_name "$tag")/$(safe_name "$pod")" mkdir -p "$pod_dir" - printf '%s\t%s\n' "$pod" "$tag" > "${pod_dir}/pod.tsv" + printf '%s\t%s\t%s\n' "$pod" "$tag" "$student_names" > "${pod_dir}/pod.tsv" log "Harvesting Claude Code logs from ${pod} (${tag})" if kubectl -n "$NAMESPACE" exec "$pod" -- sh -lc ' - cd /root || exit 0 + cd / || exit 0 set -- - [ -d .claude/projects ] && set -- "$@" .claude/projects - [ -d .claude/todos ] && set -- "$@" .claude/todos - [ -f .claude.json ] && set -- "$@" .claude.json + [ -d root/.claude/projects ] && set -- "$@" root/.claude/projects + [ -d root/.claude/todos ] && set -- "$@" root/.claude/todos + [ -f root/.claude.json ] && set -- "$@" root/.claude.json + for home in workspace/home-*; do + [ -d "$home" ] || continue + [ -d "$home/.claude/projects" ] && set -- "$@" "$home/.claude/projects" + [ -d "$home/.claude/todos" ] && set -- "$@" "$home/.claude/todos" + [ -f "$home/.claude.json" ] && set -- "$@" "$home/.claude.json" + done [ "$#" -gt 0 ] || exit 0 tar -czf - "$@" ' > "${pod_dir}/claude-code-logs.tgz" 2> "${pod_dir}/harvest.err"; then @@ -285,7 +291,13 @@ import sys data = json.load(sys.stdin) for item in data.get("items", []): meta = item.get("metadata", {}) - print("{}\t{}".format(meta.get("name", ""), meta.get("labels", {}).get("research-tag", ""))) + labels = meta.get("labels", {}) + annotations = meta.get("annotations", {}) + print("{}\t{}\t{}".format( + meta.get("name", ""), + labels.get("research-tag", ""), + annotations.get("senpai/student-names", labels.get("student", "")), + )) ' > "${dest}/pod_list.tsv" pod_count="$(wc -l < "${dest}/pod_list.tsv" | tr -d ' ')" @@ -296,14 +308,14 @@ for item in data.get("items", []): fi copy_failed=0 - while IFS="$(printf '\t')" read -r pod tag; do + while IFS="$(printf '\t')" read -r pod tag student_names; do [ -n "$pod" ] || continue while true; do job_count="$(jobs -pr | wc -l | tr -d ' ')" [ "$job_count" -lt "$MAX_PARALLEL_COPIES" ] && break sleep 1 done - copy_one_pod "$pod" "$tag" "$dest" & + copy_one_pod "$pod" "$tag" "$student_names" "$dest" & done < "${dest}/pod_list.tsv" for job in $(jobs -pr); do diff --git a/senpai.yaml b/senpai.yaml index 36100dca0..85f963992 100644 --- a/senpai.yaml +++ b/senpai.yaml @@ -51,7 +51,9 @@ student_assignment_drift_grace_s: 1800 # students n_students: 4 student_prefix: "" +# In packed mode, this is the GPU request for the shared student pod. gpus_per_student: 8 +students_per_gpu_pod: 1 cpu_per_gpu: 15 memory_gi_per_gpu: 120 preflight_only: false