From f15d13370d824a2ac91bcd4c2f8d61930ae01b7d Mon Sep 17 00:00:00 2001 From: John Wu Date: Thu, 22 Oct 2026 16:29:34 -0500 Subject: [PATCH] Add PyHealth training environments, local vLLM backend, and script runner - Pluggable LLM backends (backends.py): Gemini and OpenAI-compatible servers (vLLM), selectable per student; google-genai is now an optional extra. Non-retryable API errors (e.g. context overflow) end the student's attempt instead of crashing the run; Qwen-style reasoning blocks are stripped. - TaskEnvironment seam (environments.py) threaded through the learning cycle and final evaluation; the upstream DataFrame path is unchanged. - PyHealth environments (pyhealth_env/): students write solution(train_dataset, val_dataset, device) -> model; each attempt trains in a fresh subprocess on a GPU slot (GpuPool) and is scored on a held-out patient-level split. Built-in MIMIC-III/IV and eICU specs, reuse of cached SampleDatasets, prompts with a real-batch schema table plus PyHealth Trainer and custom BaseModel templates. - Headless runner (scripts/run_classroom.py, run.py, configs/): YAML configs with dotted overrides; notebook/Colab/Plotly code removed. - pyproject: repro extra pinning the classroom env used for reported results, separate vllm-server extra, explicit package discovery. - README: two-env setup, install pitfalls, vLLM recipes for RTX 2080 Ti and A6000, and reference MIMIC-IV mortality results (Qwen3.5-9B-AWQ students: test AUROC 0.785-0.792 vs 0.781 for the PyHealth RNN baseline). Co-Authored-By: Claude Opus 5.5 --- .env.example | 11 +- .gitignore | 4 +- README.md | 451 +++++++++++------- codegen_classroom/__init__.py | 4 + codegen_classroom/agent_config.py | 46 +- codegen_classroom/agents.py | 90 ++-- codegen_classroom/agents_test.py | 10 - codegen_classroom/backends.py | 218 +++++++++ codegen_classroom/backends_test.py | 194 ++++++++ codegen_classroom/code_execution.py | 132 ++++- codegen_classroom/code_execution_test.py | 76 +++ codegen_classroom/data_utils.py | 4 +- codegen_classroom/data_utils_test.py | 16 +- codegen_classroom/environments.py | 77 +++ codegen_classroom/metrics.py | 8 + codegen_classroom/plotting.py | 430 ++--------------- codegen_classroom/pyhealth_env/__init__.py | 48 ++ codegen_classroom/pyhealth_env/environment.py | 127 +++++ codegen_classroom/pyhealth_env/prepare.py | 296 ++++++++++++ codegen_classroom/pyhealth_env/prompts.py | 311 ++++++++++++ .../pyhealth_env/pyhealth_env_test.py | 295 ++++++++++++ codegen_classroom/pyhealth_env/registry.py | 139 ++++++ codegen_classroom/pyhealth_env/runner.py | 199 ++++++++ codegen_classroom/pyhealth_env/spec.py | 88 ++++ codegen_classroom/reporting.py | 152 +----- codegen_classroom/run.py | 425 +++++++++++++++++ codegen_classroom/run_test.py | 105 ++++ codegen_classroom/student_learning.py | 128 +++-- codegen_classroom/student_learning_test.py | 6 +- configs/mimic4_mortality_cached_qwen.yaml | 55 +++ configs/mimic4_mortality_mixed.yaml | 20 + configs/mimic4_mortality_qwen9b_2080ti.yaml | 53 ++ configs/mimic4_mortality_vllm.yaml | 45 ++ configs/regression_toy.yaml | 25 + configs/run_classroom.sub | 22 + pyproject.toml | 62 ++- scripts/run_classroom.py | 70 +++ scripts/serve_vllm.sh | 47 ++ 38 files changed, 3661 insertions(+), 828 deletions(-) create mode 100644 codegen_classroom/backends.py create mode 100644 codegen_classroom/backends_test.py create mode 100644 codegen_classroom/environments.py create mode 100644 codegen_classroom/pyhealth_env/__init__.py create mode 100644 codegen_classroom/pyhealth_env/environment.py create mode 100644 codegen_classroom/pyhealth_env/prepare.py create mode 100644 codegen_classroom/pyhealth_env/prompts.py create mode 100644 codegen_classroom/pyhealth_env/pyhealth_env_test.py create mode 100644 codegen_classroom/pyhealth_env/registry.py create mode 100644 codegen_classroom/pyhealth_env/runner.py create mode 100644 codegen_classroom/pyhealth_env/spec.py create mode 100644 codegen_classroom/run.py create mode 100644 codegen_classroom/run_test.py create mode 100644 configs/mimic4_mortality_cached_qwen.yaml create mode 100644 configs/mimic4_mortality_mixed.yaml create mode 100644 configs/mimic4_mortality_qwen9b_2080ti.yaml create mode 100644 configs/mimic4_mortality_vllm.yaml create mode 100644 configs/regression_toy.yaml create mode 100644 configs/run_classroom.sub create mode 100755 scripts/run_classroom.py create mode 100755 scripts/serve_vllm.sh diff --git a/.env.example b/.env.example index 97838fb..4c2db61 100644 --- a/.env.example +++ b/.env.example @@ -1,7 +1,16 @@ # Example Environment Variables for Codegen Classroom # Copy this file to .env and fill in your values. -# Gemini API Key (Required for student agents) +# Default LLM backend: "gemini" or "openai_compat" (alias "vllm"). +# Run configs (configs/*.yaml) can set llm.backend per run or per student. +# CODEGEN_BACKEND=vllm + +# OpenAI-compatible server (vLLM) URL, see scripts/serve_vllm.sh. +# VLLM_BASE_URL=http://localhost:8000/v1 +# Only needed if the server was started with --api-key. +# VLLM_API_KEY= + +# Gemini API Key (only for students using the gemini backend) GEMINI_API_KEY=your_api_key_here # Optional: Default output directory for experiment artifacts diff --git a/.gitignore b/.gitignore index 95c839e..02b7a09 100644 --- a/.gitignore +++ b/.gitignore @@ -6,4 +6,6 @@ .pytest_cache/ .venv/ __pycache__/ -data/ \ No newline at end of file +data/ +runs/ +logs/ diff --git a/README.md b/README.md index cd59db1..82f1ba5 100644 --- a/README.md +++ b/README.md @@ -1,202 +1,331 @@ -# Codegen Classroom 🎓 +# Codegen Classroom 🎓 (PyHealth fork) -### A framework for iterative code generation for ML and data analysis tasks -### using parallel LLM agents in Colab. +### Iterative code generation for ML tasks with parallel LLM agents, run as plain Python scripts. -------------------------------------------------------------------------------- ## Overview Codegen Classroom **automates code generation** for tasks that can be scored -(e.g., regression, classification, signal processing) using **parallel -[Gemini](https://ai.google.dev/gemini-api/docs/models) agents** in Colab. The -library operates as an **autonomous LLM agent loop**; it -evaluates the output of a language model iteratively. It is best suited for -tasks where solution quality can be **scored numerically**. Agents -**iteratively write, test, and refine** solutions in a learning loop where code -is **executed and scored** against quantitative metrics. - -Agents **improve based on performance** feedback (scores, errors) and can -**collaborate** with other agents to share insights and escape local optima. - -This framework elevates the human role from **developer to supervisor**. You -define the **task, data, constraints, and scoring** criteria, while agents -perform the **implementation and iterative refinement**. +numerically. A "classroom" of LLM **students** writes solutions, which are +**executed and scored**; students then **refine** their code from scores, +tracebacks and execution output, and **learn from peers** (mutation / +crossover) when they stall. + +This fork adds: + +* **PyHealth training environments**: students write PyTorch / PyHealth model + code for a clinical dataset + task (e.g. MIMIC-IV mortality). Each attempt + trains in a fresh subprocess on an assigned GPU and is scored with + AUROC/AUPRC (or accuracy/F1 for multiclass) on a held-out split. +* **Local LLM inference**: students can be served by a local **vLLM** server + (any OpenAI-compatible endpoint) as well as the Gemini API, including mixed + classrooms. +* **No notebooks**: experiments run from `scripts/run_classroom.py` with a + YAML config; logs, results, `report.md` and plots go to a run directory. -------------------------------------------------------------------------------- -## Getting Started +## Setup -### Prerequisites +Two separate environments are required. vLLM 0.31 ships torch 2.13 (CUDA 13), +which conflicts with PyHealth's torch 2.7.1, so the **LLM server** and the +**classroom** (which trains PyHealth models) must never share an env. -* Python 3.9 or higher -* Google Colab or Jupyter Notebook environment +```bash +git clone https://github.com/jhnwu3/codegen-classroom-pyhealth.git +cd codegen-classroom-pyhealth -### Installation & Setup +# 1. Classroom env: Python 3.12 with PyHealth installed (editable checkout, +# see "Reproducing our results" for the exact commit). +conda activate PyHealthAgent +python -m pip install -e ".[repro,cli]" # exact pinned versions we used +# or, unpinned: python -m pip install -e ".[openai,cli,dev,pyhealth]" -We recommend using a virtual environment to avoid dependency conflicts. +# 2. LLM server env (separate!) +conda create -n vllm -c conda-forge --override-channels python=3.12 -y +TMPDIR=$HOME/pip_tmp ~/miniconda3/envs/vllm/bin/python -m pip install \ + vllm==0.31.0 bitsandbytes==0.50.2 +``` -1. **Create a virtual environment**: - ```bash - python3 -m venv .venv - ``` +Pitfalls we hit on the lab machines: -2. **Activate the virtual environment**: - * On Linux/macOS: - ```bash - source .venv/bin/activate - ``` - * On Windows: - ```bash - .venv\Scripts\activate - ``` - -3. **Install dependencies in editable mode**: - ```bash - # Standard installation (core library - matplotlib included) - pip install -e . +* Call each env's interpreter explicitly (`/bin/python -m pip`). + `conda run -n vllm pip install vllm` resolved `pip` from another env and + installed vLLM's dependencies into the classroom env, breaking PyHealth's + torch/triton/pydantic. +* `/tmp` is only 8 GB on the sunlab servers; point `TMPDIR` at a larger disk + when installing vLLM (multi-GB CUDA wheels) or pip fails with + `No space left on device`. +* The lab conda has no default channel: use `-c conda-forge --override-channels`. +* PyHealth pins `pydantic~=2.11.7`; newer `google-genai` releases need + pydantic>=2.12, so the `repro` extra pins `google-genai==2.8.0`. - # For interactive notebooks (Plotly and IPython): - pip install -e ".[notebook]" # or pip install codegen_classroom[notebook] +Optional: copy `.env.example` to `.env` to set `GEMINI_API_KEY`, +`CODEGEN_BACKEND` or `VLLM_BASE_URL`. - # With LightGBM (recommended for tabular tasks): - pip install -e ".[lightgbm]" +Run the tests (366 tests, ~2 min, no GPU or LLM needed) from the repository +root: `CUDA_VISIBLE_DEVICES="" python -m pytest -q`. - # With PyTorch: - pip install -e ".[torch]" +-------------------------------------------------------------------------------- - # With JAX: - pip install -e ".[jax]" +## Running - # All dependencies (dev tools, torch, jax, lightgbm): - pip install -e ".[all]" - ``` - *Note*: `matplotlib` is a core dependency, while `plotly` and `ipython` live - in the `[notebook]` extra (`pip install codegen_classroom[notebook]`). - `lightgbm` is the approved gradient boosting framework; `xgboost` is not - supported. +1. **Start the LLM server** (with the `vllm` env's binaries on `PATH`, on + the GPUs you reserve for it; exact recipes per machine are in + [Reproducing our results](#reproducing-our-results)): -4. **Configure Environment Variables**: - Copy `.env.example` to `.env` and fill in your `GEMINI_API_KEY`. ```bash - cp .env.example .env + PATH=~/miniconda3/envs/vllm/bin:$PATH MODEL= GPUS=0 \ + bash scripts/serve_vllm.sh + curl http://localhost:8000/v1/models # check it is up ``` -5. **Run Tests**: - Verify the local installation by running the test suite: +2. **Smoke test** the pipeline on a toy regression task (no PyHealth data): + ```bash - pytest + python scripts/run_classroom.py --config configs/regression_toy.yaml ``` -6. **Set up the Colab/Jupyter Kernel** (optional): +3. **Run a PyHealth environment**: + ```bash - pip install ipykernel - python -m ipykernel install --user --name=codegen-classroom --display-name "Python (Codegen Classroom)" + python scripts/run_classroom.py --list-envs + python scripts/run_classroom.py --config configs/mimic4_mortality_vllm.yaml + + # Quick run on the MIMIC-IV dev subset with 2 cycles on GPU 1 only + python scripts/run_classroom.py --config configs/mimic4_mortality_vllm.yaml \ + --override run.num_cycles=2 'task.gpus=[1]' \ + task.overrides.dataset_kwargs.dev=true ``` + `--override key.path=value` accepts any config field (values are parsed as + YAML); `--print-config` shows the resolved config. + `configs/run_classroom.sub` is an HTCondor submit example. + +Each run writes `/_/` containing +`config.yaml`, `run.log`, per-student cycle CSVs (`students/`), every +execution's code, stdout/stderr and predictions (`executions/`), +`final_eval_results.json`, `best_analysis.json`, `report.md` and plots. +The first run of an environment also prepares and caches the PyHealth +SampleDataset and patient-level splits under `envs//`. + +-------------------------------------------------------------------------------- + +## Reproducing our results + +### Software + +| Component | Version used | +|---|---| +| Classroom env | Python 3.12.12, `pip install -e ".[repro,cli]"` (all pins in `pyproject.toml`) | +| PyHealth | editable install of [evanfebrianto/PyHealth](https://github.com/evanfebrianto/PyHealth) branch `feat/ehrmamba`, commit `c232681` (reports 2.0.1). The checkout had local edits only in FHIR-related files and temp-dir cleanup, none on the code paths used here | +| vLLM env | Python 3.12.15, vllm 0.31.0, torch 2.13.0 (CUDA 13), bitsandbytes 0.50.2 | +| NVIDIA driver | 595.71 (CUDA 13.2) on sunlab-serv-02 | + +```bash +git clone https://github.com/evanfebrianto/PyHealth.git && cd PyHealth +git checkout c232681 && python -m pip install -e . # into the classroom env +``` + +### Data + +MIMIC-IV v2.2 mortality (`pyhealth.tasks.MortalityPredictionMIMIC4`, tables +patients, admissions, diagnoses_icd, procedures_icd, prescriptions). We reuse +an existing PyHealth task cache instead of rebuilding it: + +``` +/shared/eng/pyhealth_agent/cache/b77effea-2e99-5b1b-90c3-0fa6d2540af1/tasks/ + MortalityPredictionMIMIC4_187a839c-4f67-585a-bbb3-355429e27594/ + samples_cdbbc602-34e2-5a41-8643-4c76b08829f6.ld +``` + +128,236 visits from 57,701 patients; inputs `conditions` (vocab 18,356), +`procedures` (9,781), `drugs` (6,071); ~2.3% positive. The env splits patients +80/10/10 with seed 42 (102,276 / 12,990 / 12,970 samples) and caps learning +cycles at 20,000 training samples. Without the cache, set +`task.overrides.dataset_kwargs.ehr_root` to a MIMIC-IV 2.2 root and drop +`sample_dataset_path` to rebuild it (slow, and writes to `cache_dir`). + +### LLM checkpoints + +| Model | HF revision | Size | Where it fits | +|---|---|---|---| +| [QuantTrio/Qwen3.5-9B-AWQ](https://huggingface.co/QuantTrio/Qwen3.5-9B-AWQ) | `938f8e3` | 12.4 GB | 2× RTX 2080 Ti (TP=2) | +| [unsloth/Qwen3.8-27B-unsloth-bnb-4bit](https://huggingface.co/unsloth/Qwen3.8-27B-unsloth-bnb-4bit) | `8aa5f05` | 21 GB | one free A6000 (~25-30 GB); not yet run | + +Download: `hf download --revision `. + +### Run on 8× RTX 2080 Ti (sunlab-serv-02) + +```bash +# Terminal 1: LLM on GPUs 0-1 (startup ~6 min, mostly a ~5 min warmup) +PATH=~/miniconda3/envs/vllm/bin:$PATH MODEL=QuantTrio/Qwen3.5-9B-AWQ \ + GPUS=0,1 TP=2 MAX_MODEL_LEN=32768 \ + EXTRA_ARGS='--dtype half --max-num-seqs 16 --limit-mm-per-prompt {"image":0,"video":0}' \ + bash scripts/serve_vllm.sh + +# Terminal 2: classroom, students train on GPUs 2-3 +conda activate PyHealthAgent +python scripts/run_classroom.py --config configs/mimic4_mortality_qwen9b_2080ti.yaml +``` + +Pick any two idle GPUs for the server (`GPUS=...`) and point `task.gpus` at +others (e.g. `--override 'task.gpus=[0, 1]'`). Student models need only 1-3 GB, +so they can share a card with other light jobs; the server needs whole cards +(`GPU_MEM_UTIL=0.90` of 11 GB). The first start on a machine spends ~5 min on +warmup; later starts reuse the compile cache in `~/.cache/vllm` and are ready +in about a minute. + +Turing (sm75) specifics: no bf16 (`--dtype half`); the default +`--max-num-seqs 256` exceeds Qwen3.5's linear-attention cache blocks and vLLM +refuses to start; a 16k context overflowed with within-cycle retry history +(use 32k: the KV cache holds ~155k tokens); thinking is disabled in the config +to save context. Throughput is ~70-90 tokens/s per request. + +### Run on A6000s (sunlab-serv-03) + +Use `configs/mimic4_mortality_cached_qwen.yaml` (4-bit Qwen3.8-27B) with one +free A6000 for the server; set `task.gpus` to the student GPUs. On +2026-10-16 every GPU there had only ~7 GB free (resident SGLang servers), so +this configuration has not been run yet. + +### Reference results (held-out test split, MIMIC-IV mortality) + +Hand-written templates through the same runner (no LLM), full training split, +5 epochs, RTX A6000: + +| Solution | AUROC | AUPRC | Train time | +|---|---|---|---| +| Template A: PyHealth `RNN` + `Trainer` (the seed solution) | 0.781 | 0.088 | 241 s | +| Template B: custom `BaseModel` (bag-of-codes MLP) | 0.774 | 0.078 | 91 s | + +LLM classroom runs (`configs/mimic4_mortality_qwen9b_2080ti.yaml`: 2 students, +2 cycles, no seed solution, RTX 2080 Ti): + +| Run | Students | Best solution | Val AUROC (cycle 0 → 1) | Test AUROC | Test AUPRC | +|---|---|---|---|---|---| +| `mimic4_mortality_20261019_161455` | Qwen3.5-9B-AWQ ×2 | student 1, cycle 1: custom `BaseModel` with learned attention-weighted pooling per code type, PyHealth `Trainer`, 5 epochs | 0.614 → 0.683 | **0.792** | **0.097** | +| `mimic4_mortality_20261022_155823` (rerun of the recipe above, `task.gpus=[0, 1]`, LLM on GPUs 2,6) | Qwen3.5-9B-AWQ ×2 | student 2, cycle 1: custom `BaseModel`, mean pooling + MLP | 0.614 → 0.699 | **0.785** | **0.090** | + +Notes on this run: +* Each run is a single seed with only 2 cycles; treat them as smoke tests, + not a benchmark. Across the two runs the best test AUROC was 0.785-0.792, + and in both runs one student was dropped after cycle 0. +* Student 2 failed all three cycle-0 attempts with the same error (batch + tensors not moved to `model.device`). As in upstream, a student whose cycle + fails is removed from the classroom, so student 2 took no further part. +* Common 9B-model errors, each fixed from traceback feedback within a cycle: + using `self.label_key` without setting it, label/logit shape mismatches, + missing imports, and CPU/GPU device mismatches. + +Run artifacts (code, logs and predictions for every attempt) are written to +`runs//`, and the best result of each run is appended to +`runs/leaderboard.tsv`. + +-------------------------------------------------------------------------------- + +## Configuration + +Key fields (see `codegen_classroom/run.py:DEFAULT_CONFIG` for all of them): + +| Field | Meaning | +|---|---| +| `llm.backend` | `vllm` / `openai_compat` or `gemini` (default for all students) | +| `llm.base_url` | OpenAI-compatible server URL (default `http://localhost:8000/v1`) | +| `llm.models` | One student per entry: a model name, or `{backend, model, base_url}` to mix backends | +| `llm.extra_body` | Extra request fields for vLLM (e.g. `top_p`, `chat_template_kwargs`) | +| `run.num_cycles` | Learning cycles before the final evaluation | +| `classroom.seed_solution` | Start `first`/`all`/`none` students from the PyHealth RNN baseline | +| `task.env` | Built-in PyHealth env (`--list-envs`) or `custom` | +| `task.overrides` | `EnvSpec` overrides: dataset paths/kwargs, `max_learning_train_samples`, timeouts, `metric_weights`, ... | +| `task.gpus`, `task.jobs_per_gpu` | GPUs used to train student models, and concurrent jobs per GPU | + +**GPU layout.** Keep the vLLM server and student training on different GPUs, +e.g. on 3× A6000: GPU 0 serves the LLM, and `task.gpus: [1, 2]` trains student +models. Students run concurrently; when every GPU slot is busy, the remaining +students wait for one. + -------------------------------------------------------------------------------- -## Quickstart +## PyHealth environments -The core library provides a modular agent loop that can be imported and -customized for any scorable machine learning or data science task. +Students write: ```python -import os -from codegen_classroom import agent_config -from codegen_classroom import agents -from codegen_classroom import student_learning -import pandas as pd - -# 1. Prepare training, validation, and held-out test splits -x_train = pd.DataFrame({"x1": [1.0, 2.0, 3.0, 4.0], "x2": [10.0, 20.0, 30.0, 40.0]}) -y_train = pd.DataFrame({"target": [2.1, 4.2, 5.9, 8.1]}) -x_val = pd.DataFrame({"x1": [5.0, 6.0], "x2": [50.0, 60.0]}) -y_val = pd.DataFrame({"target": [10.2, 12.1]}) -x_test = pd.DataFrame({"x1": [7.0], "x2": [70.0]}) -y_test = pd.DataFrame({"target": [14.0]}) - -# 2. Configure student agents -student_configs = student_learning.get_classroom_student_configs( - system_instruction="You are an expert ML engineer.", - model_names=["gemini-2.5-flash"], - api_key=os.environ.get("GEMINI_API_KEY"), -) -classroom = [agents.StudentAgent(cfg) for cfg in student_configs] - -# 3. Run iterative learning cycles -prompt = """ -Write a python function `solution(x_train, y_train, x_val)` that trains a regression -model on (x_train, y_train) and returns predictions on x_val as a pandas DataFrame with column 'target'. -""" -for cycle in range(3): - results = student_learning.run_parallel_classroom_learning_cycle( - current_classroom=classroom, - x_train=x_train, - y_train=y_train, - x_val=x_val, - y_val=y_val, - current_cycle=cycle, - student_initial_prompt=prompt, - student_function_name="solution", - metric_weights={"r2": 1.0}, - early_stopping_patience=2, - student_refine_instruction_options=["Consider feature scaling or Ridge regression."], - max_mutation_count=2, - mutation_strategy_options=[agents.MutationStrategy.SELF_BEST_CODE], - code_execution_max_retries=1, - ) - -# 4. Final evaluation on held-out test data -eval_results, _ = student_learning.run_parallel_classroom_final_evaluation( - students={s.name: s for s in classroom}, - x_learning=x_train, - y_learning=y_train, - x_test_final=x_test, - y_test_final=y_test, - student_function_name="solution", - metric_weights={"r2": 1.0}, -) -print("Final evaluation completed:", eval_results) +def solution(train_dataset, val_dataset, device): + ... # build and train a model, e.g. pyhealth.models.RNN + Trainer + return model +``` + +* `train_dataset` / `val_dataset` are PyHealth `SampleDataset` subsets split + by patient. During learning cycles the training split can be capped + (`max_learning_train_samples`) to keep iterations fast. +* The returned model is evaluated by the runner, which calls + `model(**batch)` on the evaluation split and reads `"y_prob"`. Students + never receive the evaluation split, and `y_true` comes from the same + batches as `y_prob`, so predictions are always aligned. +* Learning cycles are scored on the validation split. The final evaluation + retrains the chosen code (`first` / `last` / `best` cycle) on the full + training split and scores it on the held-out **test** split. +* Every attempt runs as `python -m codegen_classroom.pyhealth_env.runner` in + its own process group, with `CUDA_VISIBLE_DEVICES` set to the assigned GPU + and a wall-clock timeout (`learning_timeout_seconds`, + `final_timeout_seconds`). Errors and timeouts are fed back to the student. + +Built-in environments mirror `benchmark_tasks/` (`mimic4_mortality`, +`mimic4_readmission`, `mimic4_los`, `mimic3_mortality`, `eicu_mortality`) +with the cluster's dataset roots and cache dirs as defaults. A custom env needs +either a `dataset_class` + `task_class` (fully qualified names and kwargs) or a +prebuilt `sample_dataset_path`: + +```yaml +task: + type: pyhealth + env: custom + overrides: + name: my_task + dataset_class: pyhealth.datasets.MIMIC4Dataset + dataset_kwargs: {ehr_root: /path/to/mimiciv/2.2, ehr_tables: [diagnoses_icd], cache_dir: /path/cache} + task_class: pyhealth.tasks.MortalityPredictionMIMIC4 + label_key: mortality + task_type: binary ``` -------------------------------------------------------------------------------- ## Library Structure -The core library is organized into the following modules: - -* [agent_config.py](./codegen_classroom/agent_config.py): - Global configuration defaults for models, generation parameters, and - execution timeouts. -* [agents.py](./codegen_classroom/agents.py): Defines the `StudentAgent`, - including **state management** (`best_score`, `best_solution`), - **refinement logic**, and **mutation strategies** (self-reflection, peer - crossover). -* [code_execution.py](./codegen_classroom/code_execution.py): Handles - **extraction of Python code** blocks from LLM responses and **executes the - code** in a sandboxed environment, capturing `stdout`, `stderr`, and - results. -* [data_utils.py](./codegen_classroom/data_utils.py): Helpers for **data - loading**, generic file operations via `etils.epath`, **run ID - generation**, and local leaderboard tracking. -* [metrics.py](./codegen_classroom/metrics.py): Contains definitions for - individual **performance metrics** (e.g., MAE, R², Pearson, Worst-Group - Accuracy), including computation and **normalization logic**. -* [plotting.py](./codegen_classroom/plotting.py): Progress and final - evaluation charts using Matplotlib (headless image export) or Plotly - (interactive notebooks). -* [reporting.py](./codegen_classroom/reporting.py): Centralized logging setup, - markdown table rendering with configurable decimal precision, - `ReportSection` structured models, and dual sinks: interactive inline - display and disk export (`report.md`). -* [scoring.py](./codegen_classroom/scoring.py): Implements **solution - validation** (e.g., checking prediction shapes, NaNs) and calculates the - overall **weighted score** based on metrics. -* [student_learning.py](./codegen_classroom/student_learning.py): - **Orchestrates** the entire agent learning process. Contains the main - `run_learning_cycle` function and helpers to **run agents in parallel** and - perform final evaluations across strategies (`first`, `last`, `best`). +* [agents.py](./codegen_classroom/agents.py): `StudentAgent` (state, + refinement, mutation strategies) and `SolutionAnalysisAgent`, built on a + backend-agnostic `LLMAgent`. +* [backends.py](./codegen_classroom/backends.py): `GeminiBackend` and + `OpenAICompatBackend` (vLLM), plus the retryable-error definitions. +* [agent_config.py](./codegen_classroom/agent_config.py): defaults (backend, + vLLM URL, models, generation parameters, timeouts) and `.env` loading. +* [student_learning.py](./codegen_classroom/student_learning.py): the + learning cycle, parallel classroom orchestration and final evaluation. +* [environments.py](./codegen_classroom/environments.py): the + `TaskEnvironment` interface, which lets tasks run student code their own + way instead of the default in-process DataFrame contract. +* [pyhealth_env/](./codegen_classroom/pyhealth_env/): PyHealth env specs and + registry, data preparation, the subprocess runner, and student prompts. +* [code_execution.py](./codegen_classroom/code_execution.py): code + extraction, the forked in-process executor (DataFrame tasks), the + subprocess executor and `GpuPool`. +* [scoring.py](./codegen_classroom/scoring.py) / + [metrics.py](./codegen_classroom/metrics.py): validation, metric registry + and weighted overall score. +* [reporting.py](./codegen_classroom/reporting.py) / + [plotting.py](./codegen_classroom/plotting.py): logging, console output, + `report.md` and Matplotlib plots. +* [run.py](./codegen_classroom/run.py): the config-driven experiment loop + behind `scripts/run_classroom.py`. + +The DataFrame API from upstream still works for tabular tasks: +`solution(x_train, y_train, x_val) -> (y_pred_df, feedback)` via +`student_learning.run_parallel_classroom_learning_cycle(...)` without `env` +(see `configs/regression_toy.yaml`). -------------------------------------------------------------------------------- @@ -204,7 +333,7 @@ The core library is organized into the following modules: **This is not an officially supported Google product.** -The agents powered by Gemini are experimental and may generate inaccurate, -hallucinated, or unoptimized code. Always execute the agent loop inside a -sandboxed environment (like Google Colab) to prevent unexpected side effects -from code execution. +LLM-generated code may be inaccurate or unoptimized. Student code is executed +on this machine: subprocesses isolate crashes, GPUs and timeouts, but they are +not a security sandbox. Run experiments on machines and accounts where that is +acceptable. diff --git a/codegen_classroom/__init__.py b/codegen_classroom/__init__.py index 3359731..066cb57 100644 --- a/codegen_classroom/__init__.py +++ b/codegen_classroom/__init__.py @@ -16,8 +16,10 @@ from codegen_classroom import agent_config from codegen_classroom import agents +from codegen_classroom import backends from codegen_classroom import code_execution from codegen_classroom import data_utils +from codegen_classroom import environments from codegen_classroom import metrics from codegen_classroom import plotting from codegen_classroom import reporting @@ -29,8 +31,10 @@ __all__ = [ "__version__", "agents", + "backends", "code_execution", "data_utils", + "environments", "agent_config", "metrics", "plotting", diff --git a/codegen_classroom/agent_config.py b/codegen_classroom/agent_config.py index 3cb7c8b..6637ed3 100644 --- a/codegen_classroom/agent_config.py +++ b/codegen_classroom/agent_config.py @@ -16,13 +16,39 @@ from __future__ import annotations -import importlib import os import dotenv from etils import epath +def load_env(env_path: str | epath.PathLike | None = None) -> None: + """Loads API keys and settings from .env files into os.environ.""" + if env_path is not None: + dotenv.load_dotenv(env_path) + dotenv.load_dotenv(epath.Path(__file__).resolve().parent / ".env") + dotenv.load_dotenv() + + +# Load .env before reading environment-driven defaults below. +load_env() + +# Default LLM backend: "gemini" (Gemini API) or "openai_compat" (a local vLLM +# server or any other OpenAI-compatible endpoint). +DEFAULT_BACKEND = os.environ.get("CODEGEN_BACKEND", "gemini") + +# Default base URL of the OpenAI-compatible (vLLM) server. +DEFAULT_VLLM_BASE_URL = os.environ.get( + "VLLM_BASE_URL", "http://localhost:8000/v1" +) + +# Example open-weight models to serve with vLLM. The served model name must +# match what `vllm serve` was started with (see scripts/serve_vllm.sh). +DEFAULT_VLLM_MODELS = ( + "unsloth/Qwen3.8-27B-unsloth-bnb-4bit", + "Qwen/Qwen3-Coder-30B-A3B-Instruct-FP8", +) + # Default model used for standard agents and generic tasks. # See available models at: https://ai.google.dev/gemini-api/docs/models DEFAULT_MODEL = "gemini-3.5-flash" @@ -69,21 +95,3 @@ # Combined tuple of default models used when spawning a classroom of students. DEFAULT_STUDENT_MODELS = DEFAULT_GEMMA_MODELS + DEFAULT_GEMINI_MODELS - -def load_env(env_path: str | epath.PathLike | None = None) -> None: - """Loads GEMINI_API_KEY from package .env or Colab userdata into os.environ.""" - if env_path is not None: - dotenv.load_dotenv(env_path) - dotenv.load_dotenv(epath.Path(__file__).resolve().parent / ".env") - dotenv.load_dotenv() - if not os.environ.get("GEMINI_API_KEY"): - try: - userdata = importlib.import_module("google.colab.userdata") - - if key := userdata.get("GEMINI_API_KEY"): - os.environ["GEMINI_API_KEY"] = key - except Exception: # pylint: disable=broad-exception-caught - pass - - -load_env() diff --git a/codegen_classroom/agents.py b/codegen_classroom/agents.py index a432fbb..7c80821 100644 --- a/codegen_classroom/agents.py +++ b/codegen_classroom/agents.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Gemini agents and their configs: Student and Solution Analysis.""" +"""LLM agents and their configs: Student and Solution Analysis.""" from __future__ import annotations @@ -27,8 +27,8 @@ from typing import Any from codegen_classroom import agent_config +from codegen_classroom import backends from codegen_classroom import data_utils -from google import genai import numpy as np logger = logging.getLogger(__name__) @@ -77,10 +77,17 @@ class GenerationError(Exception): @dataclasses.dataclass class AgentConfig: - """Configuration for a Gemini/Gemma agent.""" + """Configuration for an LLM agent (Gemini API or OpenAI-compatible server).""" name: str api_key: str | None = dataclasses.field(default=None, repr=False) + # "gemini" or "openai_compat" (aliases: "openai", "vllm"). + backend: str = agent_config.DEFAULT_BACKEND + # Server URL for OpenAI-compatible backends, e.g. http://localhost:8000/v1. + base_url: str | None = None + # Extra request fields for OpenAI-compatible backends (e.g. top_p, + # chat_template_kwargs). + extra_body: Mapping[str, Any] | None = None system_instruction: str = "" model_name: str = agent_config.DEFAULT_MODEL temperature: float = agent_config.DEFAULT_TEMPERATURE @@ -123,20 +130,18 @@ def flatten(self, prefix: str = "metrics_") -> Mapping[str, Any]: return flat_dict -class GeminiAgent: - """Base class for the Student and Solution Analysis Gemini agents.""" +class LLMAgent: + """Base class for the Student and Solution Analysis LLM agents.""" def __init__(self, config: AgentConfig, run_base_dir: str): self.config = config self.name = self.config.name self.run_base_dir = run_base_dir - api_key = self.config.api_key or os.environ.get("GEMINI_API_KEY") - self.client = genai.Client( - api_key=api_key, - http_options=genai.types.HttpOptions( - timeout=self.config.http_timeout_ms - ), - ) + if backends.is_openai_compat(self.config.backend): + self.config.base_url = ( + self.config.base_url or agent_config.DEFAULT_VLLM_BASE_URL + ) + self.backend = backends.make_backend(self.config) self.system_instruction = self.config.system_instruction self.model_name = self.config.model_name self.temperature = self.config.temperature @@ -170,40 +175,27 @@ def generate(self, prompt: str) -> str: Returns: The generated response text, or an empty string if generation failed. """ - contents = [ - genai.types.Content( - role=role, - parts=[genai.types.Part.from_text(text=text)], - ) - for role, text in self.history - ] - contents.append( - genai.types.Content( - role="user", - parts=[genai.types.Part.from_text(text=prompt)], - ) + response_text = self.backend.generate( + self.history, + prompt, + system_instruction=self.system_instruction, + temperature=self.temperature, + max_output_tokens=self.max_output_tokens, + json_mode=self.config.response_mime_type == "application/json", ) - response = self.client.models.generate_content( - model=self.model_name, - contents=contents, - config=genai.types.GenerateContentConfig( - system_instruction=self.system_instruction, - max_output_tokens=self.max_output_tokens, - temperature=self.temperature, - response_mime_type=self.config.response_mime_type, - ), - ) - response_text = (response.text or "") if response else "" - if response_text.strip(): self.history.append(("user", prompt)) self.history.append(("model", response_text)) return response_text +# Backwards-compatible alias for the upstream class name. +GeminiAgent = LLMAgent + + @functools.total_ordering -class StudentAgent(GeminiAgent): +class StudentAgent(LLMAgent): """Student agent that generates solutions and learns from feedback.""" def __init__( @@ -219,7 +211,10 @@ def __init__( run_base_dir: The base directory for the current run. seed_solution: Initial model response (solution) to start with. """ - print(f"Initializing Student Agent ({config.name}) using Gemini...") + print( + f"Initializing Student Agent ({config.name}) using" + f" {config.backend}:{config.model_name}..." + ) super().__init__(config, run_base_dir) self.current_solution: str = seed_solution or "" # Full model output. self.student_dir: str = os.path.join(self.run_base_dir, "students") @@ -310,11 +305,7 @@ def generate_solution(self, prompt: str) -> str: solution = super().generate(prompt) self.current_solution = solution return solution - except ( - genai.errors.ServerError, - genai.errors.ClientError, - genai.errors.httpx.HTTPError, - ) as e: + except self.backend.retryable_exceptions as e: last_exception = e error_type = type(e).__name__ if attempt < self.config.max_code_generation_retries - 1: @@ -338,6 +329,9 @@ def generate_solution(self, prompt: str) -> str: e, exc_info=True, ) + except self.backend.fatal_exceptions as e: + logger.error("[%s] Non-retryable LLM API error: %s", self.name, e) + raise GenerationError(f"Non-retryable LLM API error: {e}") from e raise GenerationError( "Failed to generate solution after" f" {self.config.max_code_generation_retries} attempts due to API" @@ -522,7 +516,7 @@ def mutate( self.generate_solution(prompt + "\n" + final_mutation_prompt) -class SolutionAnalysisAgent(GeminiAgent): +class SolutionAnalysisAgent(LLMAgent): """Analyzes the best student model solution using an LLM agent. Extracts metadata (tags, description, insights, etc.), writes a JSON @@ -666,11 +660,7 @@ def analyze( " requested schema." ) - except ( - genai.errors.ServerError, - genai.errors.ClientError, - genai.errors.httpx.HTTPError, - ) as e: + except self.backend.retryable_exceptions as e: last_exception = e error_type = type(e).__name__ if attempt < max_retries - 1: @@ -692,6 +682,8 @@ def analyze( e, exc_info=True, ) + except self.backend.fatal_exceptions as e: + raise ValueError(f"Non-retryable LLM API error: {e}") from e raise ValueError( "Failed to obtain and parse valid JSON from model." diff --git a/codegen_classroom/agents_test.py b/codegen_classroom/agents_test.py index cd9b398..a5c8631 100644 --- a/codegen_classroom/agents_test.py +++ b/codegen_classroom/agents_test.py @@ -1009,16 +1009,6 @@ def test_load_env_with_custom_path(self): agent_config.load_env(env_path=env_file.full_path) self.assertEqual(os.environ.get('GEMINI_API_KEY'), 'custom_path_key') - def test_load_env_from_colab_userdata(self): - mock_userdata = mock.Mock() - mock_userdata.get.return_value = 'colab_test_key' - with mock.patch.dict(os.environ, {}, clear=True): - with mock.patch.dict( - sys.modules, {'google.colab.userdata': mock_userdata} - ): - agent_config.load_env() - self.assertEqual(os.environ.get('GEMINI_API_KEY'), 'colab_test_key') - if __name__ == '__main__': absltest.main() diff --git a/codegen_classroom/backends.py b/codegen_classroom/backends.py new file mode 100644 index 0000000..4ace085 --- /dev/null +++ b/codegen_classroom/backends.py @@ -0,0 +1,218 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""LLM backends: Gemini (google-genai) and OpenAI-compatible servers (vLLM). + +Each backend turns a (history, prompt) pair into response text and exposes the +tuple of exceptions that callers should treat as transient and retry on. +Provider SDKs are imported lazily so that only the selected backend's package +needs to be installed. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +import os +import re +from typing import Any, Protocol + +GEMINI = "gemini" +OPENAI_COMPAT = "openai_compat" +# Aliases accepted in configs for the OpenAI-compatible backend. +_OPENAI_COMPAT_ALIASES = frozenset({OPENAI_COMPAT, "openai", "vllm"}) + +# Reasoning blocks. Some chat templates (e.g. Qwen3.5) put the opening +# in the prompt, so the response holds only the reasoning and a closing tag. +_THINK_BLOCK_RE = re.compile(r"^.*?\s*|.*?\s*", flags=re.DOTALL) + + +def is_openai_compat(backend: str | None) -> bool: + """Returns True if `backend` names the OpenAI-compatible backend.""" + return (backend or "").lower() in _OPENAI_COMPAT_ALIASES + + +class LLMBackend(Protocol): + """Interface for a text generation backend.""" + + retryable_exceptions: tuple[type[BaseException], ...] + # Errors that retrying cannot fix (e.g. context length exceeded). Callers + # stop generating for this request instead of crashing the run. + fatal_exceptions: tuple[type[BaseException], ...] + + def generate( + self, + history: Sequence[tuple[str, str]], + prompt: str, + *, + system_instruction: str, + temperature: float, + max_output_tokens: int, + json_mode: bool = False, + ) -> str: + """Returns the model response to `prompt` given prior (role, text) turns. + + Roles in `history` are "user" and "model". + """ + ... + + +class GeminiBackend: + """Backend for the Gemini API via google-genai.""" + + def __init__( + self, + model_name: str, + api_key: str | None = None, + http_timeout_ms: int | None = None, + ): + from google import genai # pylint: disable=g-import-not-at-top + + self._genai = genai + self.model_name = model_name + self.client = genai.Client( + api_key=api_key or os.environ.get("GEMINI_API_KEY"), + http_options=genai.types.HttpOptions(timeout=http_timeout_ms), + ) + self.retryable_exceptions = ( + genai.errors.ServerError, + genai.errors.ClientError, + genai.errors.httpx.HTTPError, + ) + self.fatal_exceptions = () # All genai API errors are retried (upstream). + + def generate( + self, + history: Sequence[tuple[str, str]], + prompt: str, + *, + system_instruction: str, + temperature: float, + max_output_tokens: int, + json_mode: bool = False, + ) -> str: + types = self._genai.types + contents = [ + types.Content(role=role, parts=[types.Part.from_text(text=text)]) + for role, text in history + ] + contents.append( + types.Content(role="user", parts=[types.Part.from_text(text=prompt)]) + ) + response = self.client.models.generate_content( + model=self.model_name, + contents=contents, + config=types.GenerateContentConfig( + system_instruction=system_instruction, + max_output_tokens=max_output_tokens, + temperature=temperature, + response_mime_type="application/json" if json_mode else None, + ), + ) + return (response.text or "") if response else "" + + +class OpenAICompatBackend: + """Backend for OpenAI-compatible chat servers (vLLM, SGLang, Ollama, ...).""" + + def __init__( + self, + model_name: str, + base_url: str | None = None, + api_key: str | None = None, + http_timeout_ms: int | None = None, + extra_body: Mapping[str, Any] | None = None, + strip_think: bool = True, + ): + import openai # pylint: disable=g-import-not-at-top + + self.model_name = model_name + self.extra_body = dict(extra_body) if extra_body else None + self.strip_think = strip_think + self.client = openai.OpenAI( + base_url=base_url, + # vLLM ignores the key unless started with --api-key, but the client + # requires a non-empty value. + api_key=api_key or os.environ.get("VLLM_API_KEY") or "EMPTY", + timeout=http_timeout_ms / 1000 if http_timeout_ms else None, + max_retries=0, # Retries are handled by the agents. + ) + self.retryable_exceptions = ( + openai.APIConnectionError, # Includes APITimeoutError. + openai.RateLimitError, + openai.InternalServerError, + ) + # Remaining HTTP errors (400 context length, 401, 404, ...). Listed after + # retryable_exceptions in except clauses, so 429/5xx are still retried. + self.fatal_exceptions = (openai.APIStatusError,) + + def generate( + self, + history: Sequence[tuple[str, str]], + prompt: str, + *, + system_instruction: str, + temperature: float, + max_output_tokens: int, + json_mode: bool = False, + ) -> str: + messages = [] + if system_instruction: + messages.append({"role": "system", "content": system_instruction}) + for role, text in history: + messages.append( + {"role": "assistant" if role == "model" else role, "content": text} + ) + messages.append({"role": "user", "content": prompt}) + + kwargs: dict[str, Any] = { + "model": self.model_name, + "messages": messages, + "temperature": temperature, + "max_tokens": max_output_tokens, + } + if json_mode: + kwargs["response_format"] = {"type": "json_object"} + if self.extra_body: + kwargs["extra_body"] = self.extra_body + + response = self.client.chat.completions.create(**kwargs) + if not response or not response.choices: + return "" + text = response.choices[0].message.content or "" + if self.strip_think: + text = _THINK_BLOCK_RE.sub("", text) + return text + + +def make_backend(config: Any) -> LLMBackend: + """Creates the backend named by `config.backend` from an AgentConfig.""" + backend = (config.backend or GEMINI).lower() + if backend == GEMINI: + return GeminiBackend( + model_name=config.model_name, + api_key=config.api_key, + http_timeout_ms=config.http_timeout_ms, + ) + if is_openai_compat(backend): + return OpenAICompatBackend( + model_name=config.model_name, + base_url=config.base_url, + api_key=config.api_key, + http_timeout_ms=config.http_timeout_ms, + extra_body=config.extra_body, + ) + raise ValueError( + f"Unknown backend {config.backend!r}; expected {GEMINI!r} or" + f" {OPENAI_COMPAT!r}." + ) diff --git a/codegen_classroom/backends_test.py b/codegen_classroom/backends_test.py new file mode 100644 index 0000000..84ce6ab --- /dev/null +++ b/codegen_classroom/backends_test.py @@ -0,0 +1,194 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for LLM backends.""" + +from unittest import mock + +from absl.testing import absltest +from codegen_classroom import agents +from codegen_classroom import backends +import httpx +import openai + + +def _chat_response(text): + message = mock.Mock(content=text) + return mock.Mock(choices=[mock.Mock(message=message)]) + + +class OpenAICompatBackendTest(absltest.TestCase): + + def setUp(self): + super().setUp() + self.mock_client = mock.Mock() + self.mock_client_cls = self.enter_context( + mock.patch.object(openai, 'OpenAI', return_value=self.mock_client) + ) + self.create = self.mock_client.chat.completions.create + self.create.return_value = _chat_response('ok') + + def test_client_construction(self): + backends.OpenAICompatBackend( + model_name='m', base_url='http://host:8000/v1', http_timeout_ms=5000 + ) + _, kwargs = self.mock_client_cls.call_args + self.assertEqual(kwargs['base_url'], 'http://host:8000/v1') + self.assertEqual(kwargs['api_key'], 'EMPTY') + self.assertEqual(kwargs['timeout'], 5.0) + self.assertEqual(kwargs['max_retries'], 0) + + def test_generate_maps_roles_and_params(self): + backend = backends.OpenAICompatBackend( + model_name='qwen', extra_body={'top_p': 0.9} + ) + text = backend.generate( + [('user', 'u1'), ('model', 'm1')], + 'u2', + system_instruction='sys', + temperature=0.3, + max_output_tokens=128, + ) + self.assertEqual(text, 'ok') + _, kwargs = self.create.call_args + self.assertEqual( + kwargs['messages'], + [ + {'role': 'system', 'content': 'sys'}, + {'role': 'user', 'content': 'u1'}, + {'role': 'assistant', 'content': 'm1'}, + {'role': 'user', 'content': 'u2'}, + ], + ) + self.assertEqual(kwargs['model'], 'qwen') + self.assertEqual(kwargs['temperature'], 0.3) + self.assertEqual(kwargs['max_tokens'], 128) + self.assertEqual(kwargs['extra_body'], {'top_p': 0.9}) + self.assertNotIn('response_format', kwargs) + + def test_generate_json_mode_and_no_system(self): + backend = backends.OpenAICompatBackend(model_name='m') + backend.generate( + [], + 'p', + system_instruction='', + temperature=0.0, + max_output_tokens=8, + json_mode=True, + ) + _, kwargs = self.create.call_args + self.assertEqual(kwargs['response_format'], {'type': 'json_object'}) + self.assertEqual(kwargs['messages'], [{'role': 'user', 'content': 'p'}]) + + def test_strips_think_blocks(self): + self.create.return_value = _chat_response( + 'reasoning\nmore\nFinal answer' + ) + backend = backends.OpenAICompatBackend(model_name='m') + text = backend.generate( + [], 'p', system_instruction='', temperature=0, max_output_tokens=8 + ) + self.assertEqual(text, 'Final answer') + + def test_strips_reasoning_without_opening_tag(self): + self.create.return_value = _chat_response( + 'reasoning only\n\n\nFinal answer' + ) + backend = backends.OpenAICompatBackend(model_name='m') + text = backend.generate( + [], 'p', system_instruction='', temperature=0, max_output_tokens=8 + ) + self.assertEqual(text, 'Final answer') + + def test_empty_choices_returns_empty_string(self): + self.create.return_value = mock.Mock(choices=[]) + backend = backends.OpenAICompatBackend(model_name='m') + text = backend.generate( + [], 'p', system_instruction='', temperature=0, max_output_tokens=8 + ) + self.assertEqual(text, '') + + def test_retryable_exceptions(self): + backend = backends.OpenAICompatBackend(model_name='m') + request = httpx.Request('POST', 'http://x') + self.assertIsInstance( + openai.APITimeoutError(request), backend.retryable_exceptions + ) + self.assertIsInstance( + openai.APIConnectionError(request=request), + backend.retryable_exceptions, + ) + self.assertNotIsInstance(ValueError(), backend.retryable_exceptions) + + + def test_context_overflow_is_fatal_not_retryable(self): + backend = backends.OpenAICompatBackend(model_name='m') + response = httpx.Response(400, request=httpx.Request('POST', 'http://x')) + error = openai.BadRequestError( + 'maximum context length', response=response, body=None + ) + self.assertNotIsInstance(error, backend.retryable_exceptions) + self.assertIsInstance(error, backend.fatal_exceptions) + + +class MakeBackendTest(absltest.TestCase): + + def test_vllm_alias_and_default_base_url(self): + with mock.patch.object(openai, 'OpenAI') as mock_cls: + config = agents.AgentConfig(name='a', backend='vllm', model_name='m') + agent = agents.LLMAgent(config, self.create_tempdir().full_path) + self.assertIsInstance(agent.backend, backends.OpenAICompatBackend) + _, kwargs = mock_cls.call_args + self.assertEqual(kwargs['base_url'], 'http://localhost:8000/v1') + + def test_unknown_backend(self): + config = agents.AgentConfig(name='a', backend='nope') + with self.assertRaisesRegex(ValueError, 'Unknown backend'): + backends.make_backend(config) + + def test_student_retries_on_backend_errors(self): + with mock.patch.object(openai, 'OpenAI') as mock_cls: + create = mock_cls.return_value.chat.completions.create + create.side_effect = [ + openai.APIConnectionError(request=httpx.Request('POST', 'http://x')), + _chat_response('solution'), + ] + config = agents.AgentConfig( + name='s', + backend='openai_compat', + model_name='m', + retry_sleep_seconds=0, + ) + student = agents.StudentAgent(config, self.create_tempdir().full_path) + self.assertEqual(student.generate_solution('p'), 'solution') + self.assertEqual(create.call_count, 2) + + def test_student_fatal_error_becomes_generation_error(self): + with mock.patch.object(openai, 'OpenAI') as mock_cls: + create = mock_cls.return_value.chat.completions.create + response = httpx.Response(400, request=httpx.Request('POST', 'http://x')) + create.side_effect = openai.BadRequestError( + 'maximum context length', response=response, body=None + ) + config = agents.AgentConfig( + name='s', backend='vllm', model_name='m', retry_sleep_seconds=0 + ) + student = agents.StudentAgent(config, self.create_tempdir().full_path) + with self.assertRaisesRegex(agents.GenerationError, 'Non-retryable'): + student.generate_solution('p') + self.assertEqual(create.call_count, 1) # Not retried. + + +if __name__ == '__main__': + absltest.main() diff --git a/codegen_classroom/code_execution.py b/codegen_classroom/code_execution.py index 0039e82..e5ae213 100644 --- a/codegen_classroom/code_execution.py +++ b/codegen_classroom/code_execution.py @@ -16,7 +16,8 @@ from __future__ import annotations -from collections.abc import Mapping, Sequence +from collections.abc import Iterator, Mapping, Sequence +import contextlib import io import linecache import multiprocessing @@ -24,6 +25,7 @@ import queue import re import signal +import subprocess import sys import time from typing import Any @@ -152,11 +154,10 @@ def _execute_target( ): """Target function to run in a separate process.""" - # Prevent child from receiving SIGINT and dumping tracebacks onto Colab. + # Prevent child from receiving SIGINT and dumping tracebacks. signal.signal(signal.SIGINT, signal.SIG_IGN) - # Redirect stdout/stderr to prevent concurrent writes to Colab ZMQ sockets, - # which instantly crashes the Colab kernel. + # Redirect stdout/stderr so concurrent children do not interleave output. sys.stdout = io.StringIO() sys.stderr = io.StringIO() @@ -343,3 +344,126 @@ def execute_code( except ProcessLookupError: pass p.join() + + +# Maximum characters of child stderr included in an ExecutionError message. +MAX_SUBPROCESS_ERROR_CHARS = 8000 + + +class GpuPool: + """Thread-safe pool of GPU slots shared by concurrently executing students. + + Each device contributes `jobs_per_gpu` slots. `acquire()` blocks until a slot + is free and yields the device id to expose via CUDA_VISIBLE_DEVICES, or None + when the pool has no devices (CPU-only execution). + """ + + def __init__( + self, devices: Sequence[int | str] | None = None, jobs_per_gpu: int = 1 + ): + if jobs_per_gpu < 1: + raise ValueError("jobs_per_gpu must be >= 1.") + self.devices = [str(d) for d in (devices or [])] + self._slots: queue.Queue[str] = queue.Queue() + for _ in range(jobs_per_gpu): + for device in self.devices: + self._slots.put(device) + + @contextlib.contextmanager + def acquire(self) -> Iterator[str | None]: + if not self.devices: + yield None + return + device = self._slots.get() + try: + yield device + finally: + self._slots.put(device) + + +def _kill_process_group(proc: subprocess.Popen[Any]) -> None: + try: + os.killpg(proc.pid, signal.SIGKILL) + except (ProcessLookupError, PermissionError): + pass + try: + proc.wait(timeout=5.0) + except subprocess.TimeoutExpired: + pass + + +def execute_code_subprocess( + code: str, + command: Sequence[str], + workdir: str, + timeout_seconds: float | None = None, + env_vars: Mapping[str, str] | None = None, + code_filename: str = "solution.py", +) -> tuple[str, float]: + """Writes `code` to `workdir` and runs `command` there in a new process group. + + Unlike `execute_code`, nothing is forked from the (multithreaded) parent: the + command starts a fresh interpreter, which is safe for CUDA and lets the + child pick its GPU via CUDA_VISIBLE_DEVICES. The command is responsible for + importing `code_filename` and writing its outputs into `workdir`. + + Args: + code: The Python code to write to `workdir/code_filename`. + command: The command to run, e.g. [sys.executable, "-m", "pkg.runner", + ...]. + workdir: Directory for the code file, outputs and stdout/stderr logs. + timeout_seconds: Optional wall-clock limit; the whole process group is + killed when exceeded. + env_vars: Extra environment variables for the child. + code_filename: Name of the file the code is written to. + + Returns: + A tuple of (stdout text, execution time in seconds). + + Raises: + ExecutionError: If the process exits with a non-zero status. The message + contains the tail of stderr (including any traceback). + ExecutionTimeoutError: If the process exceeds `timeout_seconds`. + """ + os.makedirs(workdir, exist_ok=True) + with open(os.path.join(workdir, code_filename), "w") as f: + f.write(code) + stdout_path = os.path.join(workdir, "stdout.log") + stderr_path = os.path.join(workdir, "stderr.log") + env = dict(os.environ) + env.update(env_vars or {}) + + start_time = time.time() + # Log files (not pipes) avoid deadlocks on large outputs. + with open(stdout_path, "w") as out, open(stderr_path, "w") as err: + proc = subprocess.Popen( + list(command), + cwd=workdir, + stdout=out, + stderr=err, + stdin=subprocess.DEVNULL, + env=env, + start_new_session=True, # Own process group, killed as a unit. + ) + try: + returncode = proc.wait(timeout=timeout_seconds) + except subprocess.TimeoutExpired as e: + _kill_process_group(proc) + raise ExecutionTimeoutError( + f"Code execution timed out after {timeout_seconds} seconds." + ) from e + finally: + # Also covers KeyboardInterrupt and stray grandchildren (e.g. dataloader + # workers) left behind by a finished child. + _kill_process_group(proc) + exec_time = time.time() - start_time + + with open(stdout_path) as f: + stdout_text = f.read() + if returncode != 0: + with open(stderr_path) as f: + stderr_text = f.read()[-MAX_SUBPROCESS_ERROR_CHARS:] + raise ExecutionError( + f"Process exited with status {returncode}.\n{stderr_text}" + ) + return stdout_text, exec_time diff --git a/codegen_classroom/code_execution_test.py b/codegen_classroom/code_execution_test.py index cdf3138..c7f8f9c 100644 --- a/codegen_classroom/code_execution_test.py +++ b/codegen_classroom/code_execution_test.py @@ -20,7 +20,9 @@ import queue import re import signal +import sys import textwrap +import threading import time from unittest import mock @@ -738,5 +740,79 @@ def test_execute_code_background_threads_block(self, mock_mp, mock_kill): mock_kill.assert_called_with(12345, signal.SIGKILL) +class ExecuteCodeSubprocessTest(absltest.TestCase): + + def _run(self, code, timeout=None, env_vars=None): + workdir = self.create_tempdir().full_path + command = [sys.executable, 'solution.py'] + return workdir, code_execution.execute_code_subprocess( + code, command, workdir, timeout_seconds=timeout, env_vars=env_vars + ) + + def test_success_returns_stdout_and_writes_files(self): + code = textwrap.dedent(""" + import os + print('hello', os.environ['MY_VAR']) + open('out.txt', 'w').write('done') + """) + workdir, (stdout, exec_time) = self._run(code, env_vars={'MY_VAR': 'x'}) + self.assertEqual(stdout.strip(), 'hello x') + self.assertGreater(exec_time, 0) + with open(os.path.join(workdir, 'out.txt')) as f: + self.assertEqual(f.read(), 'done') + + def test_nonzero_exit_raises_with_traceback(self): + with self.assertRaisesRegex(ExecutionError, 'ZeroDivisionError'): + self._run('1 / 0') + + def test_timeout_kills_process_group(self): + code = textwrap.dedent(""" + import subprocess, sys, time + child = subprocess.Popen([sys.executable, '-c', 'import time; time.sleep(60)']) + open('child.pid', 'w').write(str(child.pid)) + time.sleep(60) + """) + workdir = self.create_tempdir().full_path + start = time.time() + with self.assertRaises(ExecutionTimeoutError): + code_execution.execute_code_subprocess( + code, [sys.executable, 'solution.py'], workdir, timeout_seconds=2 + ) + self.assertLess(time.time() - start, 20) + with open(os.path.join(workdir, 'child.pid')) as f: + child_pid = int(f.read()) + time.sleep(0.5) + with self.assertRaises(ProcessLookupError): + os.kill(child_pid, 0) + + +class GpuPoolTest(absltest.TestCase): + + def test_no_devices_yields_none(self): + pool = code_execution.GpuPool() + with pool.acquire() as device: + self.assertIsNone(device) + + def test_slots_round_robin_and_block(self): + pool = code_execution.GpuPool([1, 2], jobs_per_gpu=1) + with pool.acquire() as d1, pool.acquire() as d2: + self.assertCountEqual([d1, d2], ['1', '2']) + acquired = threading.Event() + + def _worker(): + with pool.acquire(): + acquired.set() + + t = threading.Thread(target=_worker) + t.start() + self.assertFalse(acquired.wait(timeout=0.3)) # All slots taken. + t.join(timeout=5) + self.assertTrue(acquired.is_set()) + + def test_invalid_jobs_per_gpu(self): + with self.assertRaises(ValueError): + code_execution.GpuPool([0], jobs_per_gpu=0) + + if __name__ == '__main__': absltest.main() diff --git a/codegen_classroom/data_utils.py b/codegen_classroom/data_utils.py index 07386b3..bc1617c 100644 --- a/codegen_classroom/data_utils.py +++ b/codegen_classroom/data_utils.py @@ -23,7 +23,6 @@ from typing import Any from etils import epath -from google import genai import pandas as pd logger = logging.getLogger(__name__) @@ -153,7 +152,8 @@ def _dataclass_factory(data): """Custom dict factory for dataclasses.asdict to handle ModelEntry.""" def convert_value(obj): - if hasattr(genai, "types") and isinstance(obj, genai.types.Model): + # google-genai model objects (e.g. genai.types.Model) are stringified. + if type(obj).__module__.startswith("google.genai"): return str(obj) return obj diff --git a/codegen_classroom/data_utils_test.py b/codegen_classroom/data_utils_test.py index 5eeb00d..997b3c8 100644 --- a/codegen_classroom/data_utils_test.py +++ b/codegen_classroom/data_utils_test.py @@ -222,12 +222,16 @@ def test_config_to_serializable_dict(self, cfg, expected_mutations): expected_dict.update(expected_mutations) self.assertEqual(s_dict, expected_dict) - @mock.patch('codegen_classroom.data_utils.genai') - def test_config_to_serializable_dict_with_model_entry(self, mock_genai): - """Tests conversion when config includes a ModelEntry.""" - mock_entry = mock.MagicMock() - mock_entry.__str__.return_value = "MockModel(name='test_model')" - mock_genai.types.Model = type(mock_entry) # Make isinstance work + def test_config_to_serializable_dict_with_model_entry(self): + """Tests conversion when config includes a google-genai Model entry.""" + + class Model: # Stand-in for genai.types.Model. + + def __str__(self): + return "MockModel(name='test_model')" + + Model.__module__ = 'google.genai.types' + mock_entry = Model() cfg = AnotherConfig( name='with_model', diff --git a/codegen_classroom/environments.py b/codegen_classroom/environments.py new file mode 100644 index 0000000..c85f214 --- /dev/null +++ b/codegen_classroom/environments.py @@ -0,0 +1,77 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Task environments: pluggable execution of student code for scoring. + +By default the learning loop runs student code in-process (or in a forked +child) as `solution(x_train, y_train, [aux...], x_val)` on pandas DataFrames. +A `TaskEnvironment` replaces that step for tasks whose data does not fit the +DataFrame contract (e.g. PyHealth datasets trained on GPUs). The environment +owns how the code is run and returns ground truth and predictions as +index-aligned DataFrames, so scoring, refinement and mutation are unchanged. +""" + +from __future__ import annotations + +from collections.abc import Mapping +import dataclasses +from typing import Protocol + +import pandas as pd + +# Phases passed to `TaskEnvironment.execute`. +LEARNING = "learning" # Learning cycles: scored on the validation split. +FINAL = "final" # Final evaluation: scored on the held-out test split. + + +@dataclasses.dataclass +class EnvExecution: + """Output of running student code in an environment.""" + + y_true: pd.DataFrame # Ground truth aligned to `y_pred.index`. + y_pred: pd.DataFrame # Predictions to score. + execution_feedback: str = "" # Captured stdout / feedback for the student. + exec_time: float = 0.0 # Wall-clock execution time in seconds. + auxiliary: Mapping[str, pd.DataFrame] | None = None # Extra ground truths. + + +class TaskEnvironment(Protocol): + """Runs student code for a task and returns scorable predictions.""" + + def execute( + self, + code: str, + *, + phase: str, + label: str, + timeout_seconds: float | None = None, + ) -> EnvExecution: + """Executes `code` for `phase` (LEARNING or FINAL). + + Args: + code: The extracted student code. + phase: LEARNING or FINAL. + label: Unique, filesystem-safe label for this execution (student, cycle, + attempt), used for working directories and logs. + timeout_seconds: Optional timeout override. When None, the environment + applies its own default for the phase. + + Returns: + An EnvExecution with index-aligned ground truth and predictions. + + Raises: + code_execution.ExecutionError, code_execution.ExecutionTimeoutError or + scoring.ValidatorError on failures that the student should fix. + """ + ... diff --git a/codegen_classroom/metrics.py b/codegen_classroom/metrics.py index 97200c2..39258c1 100644 --- a/codegen_classroom/metrics.py +++ b/codegen_classroom/metrics.py @@ -29,6 +29,7 @@ r2_score = sklearn.metrics.r2_score accuracy_score = sklearn.metrics.accuracy_score balanced_accuracy_score = sklearn.metrics.balanced_accuracy_score +cohen_kappa_score = sklearn.metrics.cohen_kappa_score def _correlation_compute( @@ -347,6 +348,13 @@ def normalize_score(self, raw_value: float) -> float: higher_is_better=True, normalize=None, # Range [0, 1]. Non-negative by definition. ), + 'cohen_kappa': Metric( + name='cohen_kappa', + func=cohen_kappa_score, + description="Cohen's Kappa", + higher_is_better=True, + normalize=clip_negative, # Range [-1, 1]; below chance clips to 0. + ), 'balanced_accuracy': Metric( name='balanced_accuracy', func=balanced_accuracy_score, diff --git a/codegen_classroom/plotting.py b/codegen_classroom/plotting.py index 70b899c..e59c774 100644 --- a/codegen_classroom/plotting.py +++ b/codegen_classroom/plotting.py @@ -12,7 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Plotting utils for charts, decoupled from reporting and display logic.""" +"""Plotting utils for charts, decoupled from reporting and display logic. + +Charts are rendered headlessly with Matplotlib (Agg backend) and saved to disk. +""" from __future__ import annotations @@ -29,6 +32,8 @@ import matplotlib.pyplot as plt import numpy as np +matplotlib.use("Agg") + StudentProgressData = Sequence[Mapping[str, str | int | float | bool | None]] FinalResultsData = Sequence[ # Example: [{'student_name': str, 'best': {metric: score}, 'best': ...}] @@ -58,110 +63,6 @@ def get_metric_key(metric_name: str) -> str: ) -def _apply_common_styling( - fig: Any, - width: int, - height: int, - title_font_size: int, - axis_font_size: int, - legend_font_size: int, -): - """Applies a consistent set of layout and axis styles to a Plotly figure.""" - fig.update_layout( - width=width, - height=height, - legend_font_size=legend_font_size, - ) - for annotation in fig.layout.annotations: - annotation.font.size = title_font_size - - fig.update_xaxes( - tickfont_size=axis_font_size, - title_font_size=axis_font_size, - gridcolor="lightgrey", - ) - fig.update_yaxes( - tickfont_size=axis_font_size, - title_font_size=axis_font_size, - gridcolor="lightgrey", - ) - - -def _add_cumulative_plot( - fig: Any, - student_data: list[dict[str, Any]], - higher_is_better: bool, - student_name: str, - student_color: str, - row: int, - show_mutation_points: bool = False, -): - """Adds a cumulative best trace to a Plotly subplot.""" - import plotly.graph_objects as go # pylint: disable=g-import-not-at-top - - student_data.sort(key=lambda x: x["cycle"]) - cycles = [d["cycle"] for d in student_data] - scores = [d["metric_val"] for d in student_data] - - if not cycles: - return - - cumulative_y = ( - np.maximum.accumulate(scores) - if higher_is_better - else np.minimum.accumulate(scores) - ) - cumulative_x = cycles - - mutation_x = [] - mutation_y = [] - mutation_text = [] - - if show_mutation_points: - for i, data in enumerate(student_data): - mutation_strategy = data.get("mutation_strategy", "") - if mutation_strategy and str(mutation_strategy).lower() != "nan": - mutation_x.append(cycles[i]) - mutation_y.append(cumulative_y[i]) - current_best = cumulative_y[i] - mutation_text.append( - f"Cycle: {cycles[i]}
Mutation:" - f" {mutation_strategy}
Score: {current_best:.4f}" - ) - fig.add_trace( - go.Scatter( - x=cumulative_x, - y=cumulative_y, - mode="lines+markers", - name=f"{student_name} Cumulative", - legendgroup=student_name, - showlegend=False, - line=dict(shape="hv", color=student_color), - marker=dict(size=4), - hoverinfo="skip", - ), - row=row, - col=2, - ) - - if mutation_x: - fig.add_trace( - go.Scatter( - x=mutation_x, - y=mutation_y, - mode="markers", - name=f"{student_name} Mutations", - legendgroup=student_name, - showlegend=False, - marker=dict(size=10, color=student_color, symbol="diamond"), - hovertext=mutation_text, - hoverinfo="text", - ), - row=row, - col=2, - ) - - def plot_progress( student_results_dir: str | epath.Path, metric_names: list[str], @@ -174,8 +75,7 @@ def plot_progress( legend_font_size: int = 12, show_mutation_points: bool = False, save_path: str | None = None, - show: bool = True, - backend: str = "plotly", # "plotly" or "matplotlib" + show: bool = False, ): """Plots the progress of students. @@ -191,8 +91,7 @@ def plot_progress( legend_font_size: Font size for legend items. show_mutation_points: Show mutation points on the cumulative plot. save_path: If set, save figure here (forces theme to light). - show: Whether to show the figure. - backend: "plotly" (notebook default) or "matplotlib" (fallback/write). + show: Whether to call plt.show() (no-op with the Agg backend). """ if not metric_names: print("Metrics list cannot be empty.") @@ -209,161 +108,18 @@ def plot_progress( # A committed artifact should always be light dark_mode = False - if backend == "plotly": - plot_progress_interactive( - student_data_map, - metric_names, - dark_mode, - fig_width, - fig_height_per_metric, - title_font_size, - axis_font_size, - legend_font_size, - show_mutation_points, - save_path, - show, - ) - else: - plot_progress_static( - student_data=student_data_map, - metric_names=metric_names, - dark_mode=dark_mode, - single_plot_size=(fig_width / 100.0, fig_height_per_metric / 100.0), - title_fontsize=title_font_size, - label_fontsize=axis_font_size, - legend_fontsize=legend_font_size, - show_mutation_points=show_mutation_points, - save_path=save_path, - show=show, - ) - - -def plot_progress_interactive( - student_data_map, - metric_names, - dark_mode, - fig_width, - fig_height_per_metric, - title_font_size, - axis_font_size, - legend_font_size, - show_mutation_points, - save_path, - show, -): - """Plots learning progress using the Plotly backend.""" - # lazy: notebook-only, keeps headless import light - from plotly import subplots # pylint: disable=g-import-not-at-top - import plotly.graph_objects as go # pylint: disable=g-import-not-at-top - - color_map = { - name: color - for name, color in zip( - student_data_map.keys(), - _PALETTE * (len(student_data_map) // len(_PALETTE) + 1), - ) - } - - subplot_titles = [] - valid_metrics = {} - for metric_name in metric_names: - config = metric_lib.ALL_METRICS_CONFIG.get(metric_name) - if not config: - continue - higher_is_better = config.higher_is_better - valid_metrics[metric_name] = {"higher_is_better": higher_is_better} - subplot_titles.append(f"Progress: '{metric_name}' ({config.description})") - cumulative_type = "Max" if higher_is_better else "Min" - subplot_titles.append(f"Cumulative {cumulative_type}: '{metric_name}'") - - if not valid_metrics: - return - - fig = subplots.make_subplots( - rows=len(valid_metrics), cols=2, subplot_titles=subplot_titles - ) - - for i, metric_name in enumerate(valid_metrics): - row = i + 1 - config = valid_metrics[metric_name] - has_data_for_metric = False - metric_key = get_metric_key(metric_name) - - for student_name, all_cycles_data in student_data_map.items(): - student_metrics = [] - for record in all_cycles_data: - if metric_key in record and "cycle" in record: - try: - student_metrics.append({ - "cycle": int(record["cycle"]), - "metric_val": float(record[metric_key]), - "mutation_strategy": str(record.get("mutation_strategy", "")), - }) - except (ValueError, TypeError): - continue - - if not student_metrics: - continue - - has_data_for_metric = True - student_color = color_map.get(student_name) - student_metrics.sort(key=lambda x: x["cycle"]) - cycles = [d["cycle"] for d in student_metrics] - scores = [d["metric_val"] for d in student_metrics] - - fig.add_trace( - go.Scatter( - x=cycles, - y=scores, - mode="markers", - name=student_name, - legendgroup=student_name, - marker_color=student_color, - showlegend=(i == 0), - ), - row=row, - col=1, - ) - _add_cumulative_plot( - fig, - student_metrics, - config["higher_is_better"], - student_name, - student_color or _PALETTE[0], - row, - show_mutation_points, - ) - - if not has_data_for_metric: - fig.layout.annotations[2 * i].text = f"No Data for '{metric_name}'" - fig.layout.annotations[2 * i + 1].text = f"No Data for '{metric_name}'" - - _apply_common_styling( - fig=fig, - width=fig_width, - height=fig_height_per_metric * len(valid_metrics), - title_font_size=title_font_size, - axis_font_size=axis_font_size, - legend_font_size=legend_font_size, + plot_progress_static( + student_data=student_data_map, + metric_names=metric_names, + dark_mode=dark_mode, + single_plot_size=(fig_width / 100.0, fig_height_per_metric / 100.0), + title_fontsize=title_font_size, + label_fontsize=axis_font_size, + legend_fontsize=legend_font_size, + show_mutation_points=show_mutation_points, + save_path=save_path, + show=show, ) - fig.update_xaxes(title_text="Cycle") - for i, metric_name in enumerate(valid_metrics): - display_name = metric_name.replace("_", " ").title() - fig.update_yaxes(title_text=display_name, row=i + 1, col=1) - fig.update_yaxes(title_text="Cumulative Best", row=i + 1, col=2) - - if dark_mode: - fig.update_layout(template="plotly_dark") - - if save_path: - # Requires optional dependencies kaleido or orca - try: - with epath.Path(save_path).open("wb") as f: - fig.write_image(f) - except Exception as e: # pylint: disable=broad-exception-caught - print(f"Failed to write image: {e}") - if show: - fig.show() def plot_progress_static( @@ -377,7 +133,7 @@ def plot_progress_static( show_mutation_points: bool = True, save_path: str | epath.Path | None = None, save_fig: str | epath.Path | None = None, - show: bool = True, + show: bool = False, show_fig: bool | None = None, dpi: int = 150, metric_weights: Mapping[str, float] | None = None, @@ -387,9 +143,6 @@ def plot_progress_static( if show_fig is not None: show = show_fig - if not show: - matplotlib.use("Agg") - if isinstance(student_data, (str, epath.Path)): p = epath.Path(str(student_data)) students_dir = p / "students" if (p / "students").exists() else p @@ -586,8 +339,7 @@ def plot_final_test( axis_font_size: int = 12, legend_font_size: int = 12, save_path: str | None = None, - show: bool = True, - backend: str = "plotly", # "plotly" or "matplotlib" + show: bool = False, ): """Plots final test evaluation results across First, Last, and Best strategies.""" if final_results_path: @@ -610,138 +362,18 @@ def plot_final_test( if save_path: dark_mode = False - if backend == "plotly": - _plot_final_test_plotly( - final_results, - metrics, - dark_mode, - fig_width_per_metric, - fig_height, - title_font_size, - axis_font_size, - legend_font_size, - save_path, - show, - ) - else: - _plot_final_test_matplotlib( - final_results, - metrics, - dark_mode, - (fig_width_per_metric / 100, fig_height / 100), - title_font_size, - axis_font_size, - legend_font_size, - save_path, - show, - ) - - -def _plot_final_test_plotly( - final_results, - metrics, - dark_mode, - fig_width_per_metric, - fig_height, - title_font_size, - axis_font_size, - legend_font_size, - save_path, - show, -): - """Plots final test evaluation bars using the Plotly backend.""" - # lazy: notebook-only, keeps headless import light - from plotly import subplots # pylint: disable=g-import-not-at-top - import plotly.graph_objects as go # pylint: disable=g-import-not-at-top - - student_names = [res.get("student_name", "Unknown") for res in final_results] - plot_metrics = metrics - - fig = subplots.make_subplots( - rows=1, - cols=len(plot_metrics), - subplot_titles=[ - f"Final Score: '{metric.upper()}'" for metric in plot_metrics - ], + _plot_final_test_matplotlib( + final_results, + metrics, + dark_mode, + (fig_width_per_metric / 100, fig_height / 100), + title_font_size, + axis_font_size, + legend_font_size, + save_path, + show, ) - bar_configs = { - "First": { - "color": _PALETTE[2], - "data_key": "first", - }, - "Last": { - "color": _PALETTE[0], - "data_key": "last", - }, - "Best": {"color": _PALETTE[1], "data_key": "best"}, - } - - for i, metric in enumerate(plot_metrics): - for name, config in bar_configs.items(): - data_key = config["data_key"] - metric_key = get_metric_key(metric) - - values = [] - for res in final_results: - result_dict = res.get(data_key) or {} - values.append(result_dict.get(metric_key)) - - custom_data = None - hovertemplate = ( - f"%{{x}}
Version: {name}
{metric.upper()}: %{{y}}" - ) - - if name == "Best": - custom_data = [] - for res in final_results: - best_data = res.get("best") or {} - custom_data.append(best_data.get("best_cycle")) - hovertemplate += "
Best Cycle: %{customdata}" - else: - hovertemplate += "" - - fig.add_trace( - go.Bar( - name=name, - x=student_names, - y=values, - marker_color=config["color"], - showlegend=(i == 0), - customdata=custom_data, - hovertemplate=hovertemplate, - ), - row=1, - col=i + 1, - ) - - _apply_common_styling( - fig=fig, - width=fig_width_per_metric * len(plot_metrics), - height=fig_height, - title_font_size=title_font_size, - axis_font_size=axis_font_size, - legend_font_size=legend_font_size, - ) - - fig.update_layout(barmode="group") - fig.update_xaxes(title_text="Student", tickangle=-45, automargin=True) - for i, metric in enumerate(plot_metrics): - fig.update_yaxes( - title_text=metric.replace("_", " ").upper(), row=1, col=i + 1 - ) - - if dark_mode: - fig.update_layout(template="plotly_dark") - if save_path: - try: - with epath.Path(save_path).open("wb") as f: - fig.write_image(f) - except Exception as e: # pylint: disable=broad-exception-caught - print(f"Failed to write image: {e}") - if show: - fig.show() - def _plot_final_test_matplotlib( final_results, diff --git a/codegen_classroom/pyhealth_env/__init__.py b/codegen_classroom/pyhealth_env/__init__.py new file mode 100644 index 0000000..97f2fe2 --- /dev/null +++ b/codegen_classroom/pyhealth_env/__init__.py @@ -0,0 +1,48 @@ +"""PyHealth training environments for codegen classroom. + +Students write `solution(train_dataset, val_dataset, device) -> model` for a +PyHealth dataset + task; each attempt is trained in a fresh subprocess on an +assigned GPU and scored on a held-out split. + +Typical use: + + spec = registry.get_spec("mimic4_mortality") + env, prepared = make_environment(spec, env_dir=..., work_dir=run_dir, + gpus=[1, 2]) + prompt = prompts.build_task_prompt(spec, prepared) + student_learning.run_parallel_classroom_learning_cycle(..., env=env) +""" + +from __future__ import annotations + +from collections.abc import Sequence + +from codegen_classroom import code_execution +from codegen_classroom.pyhealth_env import environment +from codegen_classroom.pyhealth_env import prepare +from codegen_classroom.pyhealth_env import prompts +from codegen_classroom.pyhealth_env import registry +from codegen_classroom.pyhealth_env import spec + +EnvSpec = spec.EnvSpec +PyHealthEnvironment = environment.PyHealthEnvironment + + +def make_environment( + env_spec: spec.EnvSpec, + env_dir: str, + work_dir: str, + gpus: Sequence[int | str] | None = None, + jobs_per_gpu: int = 1, + seed: int = 0, +) -> tuple[environment.PyHealthEnvironment, prepare.PreparedEnv]: + """Prepares the env (cached in `env_dir`) and returns it with its artifacts.""" + prepared = prepare.prepare_env(env_spec, env_dir) + env = environment.PyHealthEnvironment( + env_spec, + prepared, + work_dir=work_dir, + gpu_pool=code_execution.GpuPool(gpus, jobs_per_gpu=jobs_per_gpu), + seed=seed, + ) + return env, prepared diff --git a/codegen_classroom/pyhealth_env/environment.py b/codegen_classroom/pyhealth_env/environment.py new file mode 100644 index 0000000..c6e4e7c --- /dev/null +++ b/codegen_classroom/pyhealth_env/environment.py @@ -0,0 +1,127 @@ +"""TaskEnvironment that trains student PyHealth models in GPU subprocesses.""" + +from __future__ import annotations + +import json +import logging +import os +import re +import sys + +from codegen_classroom import code_execution +from codegen_classroom import environments +from codegen_classroom.pyhealth_env import prepare as prepare_lib +from codegen_classroom.pyhealth_env import spec as spec_lib +import numpy as np +import pandas as pd + +logger = logging.getLogger(__name__) + +RUNNER_MODULE = "codegen_classroom.pyhealth_env.runner" +# Max characters of the student's stdout returned as execution feedback. +MAX_FEEDBACK_CHARS = 8000 + + +class PyHealthEnvironment: + """Runs `solution(train_dataset, val_dataset, device) -> model` per attempt. + + Learning cycles train on the (optionally capped) train split and are scored + on the validation split. Final evaluation trains on the full train split and + is scored on the held-out test split, which students never receive. + """ + + def __init__( + self, + spec: spec_lib.EnvSpec, + prepared: prepare_lib.PreparedEnv, + work_dir: str, + gpu_pool: code_execution.GpuPool | None = None, + function_name: str = "solution", + python_executable: str = sys.executable, + seed: int = 0, + ): + self.spec = spec + self.prepared = prepared + # Absolute, because each runner subprocess starts in its own directory. + self.work_dir = os.path.abspath(work_dir) + self.gpu_pool = gpu_pool or code_execution.GpuPool() + self.function_name = function_name + self.python_executable = python_executable + self.seed = seed + + def _splits(self, phase: str) -> tuple[str, str, str]: + if phase == environments.LEARNING: + return ( + prepare_lib.LEARNING_TRAIN, + prepare_lib.LEARNING_VAL, + prepare_lib.LEARNING_VAL, + ) + if phase == environments.FINAL: + return prepare_lib.TRAIN, prepare_lib.VAL, prepare_lib.TEST + raise ValueError(f"Unknown phase {phase!r}.") + + def execute( + self, + code: str, + *, + phase: str, + label: str, + timeout_seconds: float | None = None, + ) -> environments.EnvExecution: + train_split, val_split, eval_split = self._splits(phase) + if timeout_seconds is None: + timeout_seconds = ( + self.spec.learning_timeout_seconds + if phase == environments.LEARNING + else self.spec.final_timeout_seconds + ) + safe_label = re.sub(r"[^A-Za-z0-9_.-]", "_", label) + workdir = os.path.join(self.work_dir, "executions", safe_label) + os.makedirs(workdir, exist_ok=True) + output_file = os.path.join(workdir, "predictions.npz") + if os.path.exists(output_file): + os.remove(output_file) + job = { + "sample_dataset_path": self.prepared.sample_dataset_path, + "splits_path": self.prepared.splits_path, + "train_split": train_split, + "val_split": val_split, + "eval_split": eval_split, + "label_key": self.spec.label_key, + "task_type": self.spec.task_type, + "function_name": self.function_name, + "code_file": os.path.join(workdir, "solution.py"), + "output_file": output_file, + "inference_batch_size": self.spec.inference_batch_size, + "seed": self.seed, + } + job_path = os.path.join(workdir, "job.json") + with open(job_path, "w") as f: + json.dump(job, f, indent=2) + + with self.gpu_pool.acquire() as device: + env_vars = {} + if device is not None: + env_vars["CUDA_VISIBLE_DEVICES"] = device + logger.debug("[%s] Running on GPU %s (%s phase)", label, device, phase) + stdout, exec_time = code_execution.execute_code_subprocess( + code, + [self.python_executable, "-m", RUNNER_MODULE, "--job", job_path], + workdir, + timeout_seconds=timeout_seconds, + env_vars=env_vars, + ) + + if not os.path.exists(output_file): + raise code_execution.ExecutionError( + "The runner finished without writing predictions." + ) + with np.load(output_file) as data: + y_true = pd.DataFrame({self.spec.label_key: data["y_true"]}) + y_pred = pd.DataFrame({self.spec.label_key: data["y_score"]}) + return environments.EnvExecution( + y_true=y_true, + y_pred=y_pred, + execution_feedback=stdout[-MAX_FEEDBACK_CHARS:], + exec_time=exec_time, + ) diff --git a/codegen_classroom/pyhealth_env/prepare.py b/codegen_classroom/pyhealth_env/prepare.py new file mode 100644 index 0000000..4490cd8 --- /dev/null +++ b/codegen_classroom/pyhealth_env/prepare.py @@ -0,0 +1,296 @@ +"""One-time preparation of a PyHealth environment. + +Builds (or reuses) the task's litdata SampleDataset, computes deterministic +patient-level splits and summarizes the schema and label distribution for the +student prompt. Everything is written to `env_dir` so that student executions +(separate processes) can reload the exact same splits. + +Runs on CPU in the main process (PyHealth's `set_task` refuses to run in a +non-main process). +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +import collections +import dataclasses +import hashlib +import json +import logging +import os +from typing import Any + +from codegen_classroom.pyhealth_env import spec as spec_lib +import numpy as np + +logger = logging.getLogger(__name__) + +# Bump when the prepared artifacts change so cached env dirs are rebuilt. +PREPARE_VERSION = 2 +PREPARED_FILE = "prepared.json" +SPLITS_FILE = "splits.npz" +# Max samples read to estimate the label distribution for the prompt. +_LABEL_STATS_MAX_SAMPLES = 20000 + +# Split names stored in SPLITS_FILE. +TRAIN, VAL, TEST = "train", "val", "test" +LEARNING_TRAIN, LEARNING_VAL = "learning_train", "learning_val" + + +@dataclasses.dataclass +class PreparedEnv: + """Artifacts of a prepared environment, stored in `env_dir`.""" + + env_dir: str + sample_dataset_path: str + split_sizes: Mapping[str, int] + label_distribution: Mapping[str, int] + schema_summary: Mapping[str, Any] + # Per-key description of one real batch from `get_dataloader`. + batch_summary: Mapping[str, Any] = dataclasses.field(default_factory=dict) + + @property + def splits_path(self) -> str: + return os.path.join(self.env_dir, SPLITS_FILE) + + +def _split_fingerprint(spec: spec_lib.EnvSpec) -> str: + """Hash of every spec field that affects the data or splits.""" + keys = { + "version": PREPARE_VERSION, + "dataset_class": spec.dataset_class, + "dataset_kwargs": spec.dataset_kwargs, + "task_class": spec.task_class, + "task_kwargs": spec.task_kwargs, + "sample_dataset_path": spec.sample_dataset_path, + "split_ratios": spec.split_ratios, + "split_seed": spec.split_seed, + "max_learning_train_samples": spec.max_learning_train_samples, + "max_learning_val_samples": spec.max_learning_val_samples, + "label_key": spec.label_key, + } + blob = json.dumps(keys, sort_keys=True, default=str).encode() + return hashlib.sha256(blob).hexdigest()[:16] + + +def build_sample_dataset_path(spec: spec_lib.EnvSpec) -> str: + """Returns the litdata SampleDataset directory for the spec, building it.""" + if spec.sample_dataset_path: + return spec.sample_dataset_path + dataset_cls = spec_lib.import_object(spec.dataset_class) + task_cls = spec_lib.import_object(spec.task_class) + logger.info("Loading %s(%s)", spec.dataset_class, dict(spec.dataset_kwargs)) + base_dataset = dataset_cls(**spec.dataset_kwargs) + task = task_cls(**spec.task_kwargs) + logger.info("Setting task %s (cached after the first run)", spec.task_class) + sample_dataset = base_dataset.set_task(task) + return str(sample_dataset.path) + + +def split_by_patient_indices( + patient_to_index: Mapping[Any, Sequence[int]], + ratios: Sequence[float], + seed: int, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Patient-level split into sorted sample index arrays. + + Same procedure as `pyhealth.datasets.split_by_patient`, but returns indices + so they can be persisted and reused by child processes. + + Args: + patient_to_index: Mapping from patient id to sample indices. + ratios: Train/val/test ratios summing to 1. + seed: Shuffle seed. + + Returns: + Sorted (train, val, test) sample index arrays. + """ + if not np.isclose(sum(ratios), 1.0): + raise ValueError(f"Split ratios must sum to 1, got {ratios}.") + patients = sorted(patient_to_index, key=str) + rng = np.random.default_rng(seed) + rng.shuffle(patients) + n = len(patients) + cut1 = int(n * ratios[0]) + cut2 = int(n * (ratios[0] + ratios[1])) + groups = (patients[:cut1], patients[cut1:cut2], patients[cut2:]) + + def _indices(group): + idx = [i for p in group for i in patient_to_index[p]] + return np.array(sorted(idx), dtype=np.int64) + + return tuple(_indices(g) for g in groups) + + +def _subsample( + indices: np.ndarray, max_samples: int | None, seed: int +) -> np.ndarray: + if max_samples is None or len(indices) <= max_samples: + return indices + rng = np.random.default_rng(seed) + return np.sort(rng.choice(indices, size=max_samples, replace=False)) + + +def _label_value(label: Any) -> str: + """Converts a processed label (tensor/array/scalar) to a hashable string.""" + arr = np.asarray(label.detach().cpu() if hasattr(label, "detach") else label) + if arr.size == 1: + return str(int(arr.reshape(-1)[0])) + return str(int(arr.argmax())) + + +def _processor_summary(processor: Any) -> dict[str, Any]: + summary: dict[str, Any] = {"processor": type(processor).__name__} + size = getattr(processor, "size", None) + if callable(size): + try: + summary["size"] = int(size()) + except Exception: # pylint: disable=broad-exception-caught + pass + for attr in ("schema", "is_token"): + fn = getattr(processor, attr, None) + if callable(fn): + try: + value = fn() + summary[attr] = list(value) if isinstance(value, tuple) else value + except Exception: # pylint: disable=broad-exception-caught + pass + return summary + + +def summarize_schema(sample_dataset: Any) -> dict[str, Any]: + """Summarizes input/output schemas and fitted processors for the prompt.""" + return { + "inputs": { + key: { + "schema": str(sample_dataset.input_schema.get(key)), + **_processor_summary(proc), + } + for key, proc in sample_dataset.input_processors.items() + }, + "outputs": { + key: { + "schema": str(sample_dataset.output_schema.get(key)), + **_processor_summary(proc), + } + for key, proc in sample_dataset.output_processors.items() + }, + } + + +def summarize_batch(sample_dataset: Any, batch_size: int = 4) -> dict: + """Describes each field of one real batch (type, dtype, shape).""" + from pyhealth.datasets import get_dataloader # pylint: disable=g-import-not-at-top + + batch = next(iter(get_dataloader(sample_dataset, batch_size=batch_size))) + summary = {} + for key, value in batch.items(): + if hasattr(value, "shape"): + summary[key] = { + "type": type(value).__name__, + "dtype": str(value.dtype).replace("torch.", ""), + "shape": list(value.shape), + } + elif isinstance(value, (tuple, list)) and value and hasattr( + value[0], "shape" + ): + summary[key] = { + "type": f"{type(value).__name__} of tensors", + "shapes": [list(v.shape) for v in value], + "dtypes": [str(v.dtype).replace("torch.", "") for v in value], + } + else: + summary[key] = { + "type": f"{type(value).__name__}[{type(value[0]).__name__}]" + if isinstance(value, (tuple, list)) and value + else type(value).__name__ + } + return summary + + +def prepare_env(spec: spec_lib.EnvSpec, env_dir: str) -> PreparedEnv: + """Prepares (or loads the cached) environment artifacts in `env_dir`.""" + from pyhealth.datasets import SampleDataset # pylint: disable=g-import-not-at-top + + # Absolute paths: the artifacts are read by subprocesses with another cwd. + env_dir = os.path.abspath(env_dir) + os.makedirs(env_dir, exist_ok=True) + prepared_path = os.path.join(env_dir, PREPARED_FILE) + fingerprint = _split_fingerprint(spec) + if os.path.exists(prepared_path): + with open(prepared_path) as f: + cached = json.load(f) + if cached.get("fingerprint") == fingerprint and os.path.exists( + os.path.join(env_dir, SPLITS_FILE) + ): + logger.info("Using prepared PyHealth env at %s", env_dir) + return PreparedEnv( + env_dir=env_dir, + sample_dataset_path=os.path.abspath(cached["sample_dataset_path"]), + split_sizes=cached["split_sizes"], + label_distribution=cached["label_distribution"], + schema_summary=cached["schema_summary"], + batch_summary=cached.get("batch_summary", {}), + ) + + sample_path = os.path.abspath(build_sample_dataset_path(spec)) + dataset = SampleDataset(path=sample_path) + if spec.label_key not in dataset.output_processors: + raise ValueError( + f"label_key {spec.label_key!r} not in task outputs" + f" {list(dataset.output_processors)}." + ) + + train, val, test = split_by_patient_indices( + dataset.patient_to_index, spec.split_ratios, spec.split_seed + ) + splits = { + TRAIN: train, + VAL: val, + TEST: test, + LEARNING_TRAIN: _subsample( + train, spec.max_learning_train_samples, spec.split_seed + ), + LEARNING_VAL: _subsample( + val, spec.max_learning_val_samples, spec.split_seed + 1 + ), + } + for name, idx in splits.items(): + if len(idx) == 0: + raise ValueError(f"Split {name!r} is empty; check split_ratios.") + np.savez(os.path.join(env_dir, SPLITS_FILE), **splits) + + stats_idx = _subsample(train, _LABEL_STATS_MAX_SAMPLES, spec.split_seed) + label_counts = collections.Counter( + _label_value(sample[spec.label_key]) + for sample in dataset.subset(stats_idx.tolist()) + ) + + prepared = PreparedEnv( + env_dir=env_dir, + sample_dataset_path=sample_path, + split_sizes={k: int(len(v)) for k, v in splits.items()}, + label_distribution=dict(sorted(label_counts.items())), + schema_summary=summarize_schema(dataset), + batch_summary=summarize_batch( + dataset.subset(splits[LEARNING_TRAIN][:8].tolist()) + ), + ) + with open(prepared_path, "w") as f: + json.dump( + { + "fingerprint": fingerprint, + "spec": dataclasses.asdict(spec), + **dataclasses.asdict(prepared), + }, + f, + indent=2, + default=str, + ) + logger.info( + "Prepared PyHealth env %s: splits=%s labels=%s", + spec.name, + prepared.split_sizes, + prepared.label_distribution, + ) + return prepared diff --git a/codegen_classroom/pyhealth_env/prompts.py b/codegen_classroom/pyhealth_env/prompts.py new file mode 100644 index 0000000..c79159a --- /dev/null +++ b/codegen_classroom/pyhealth_env/prompts.py @@ -0,0 +1,311 @@ +"""Student prompts for PyHealth environments.""" + +from __future__ import annotations + +import textwrap + +from codegen_classroom import scoring +from codegen_classroom.pyhealth_env import prepare as prepare_lib +from codegen_classroom.pyhealth_env import spec as spec_lib + +SYSTEM_INSTRUCTION = """\ +You are an expert clinical machine learning engineer who writes correct, +efficient PyTorch / PyHealth 2.0 code. You iteratively improve a model for a +clinical prediction task based on validation scores and execution feedback. + +Always answer with a short plan followed by ONE complete ```python code block +that defines the requested function. The code is run as a standalone module, so +include every import it needs. Do not read files outside the provided datasets, +do not download data, and do not use the network. +""" + +PYHEALTH_CHEATSHEET = """\ +PyHealth 2.0 quick reference: +- Datasets: `train_dataset` / `val_dataset` are `pyhealth.datasets.SampleDataset` + subsets. Each sample is a dict of processed tensors keyed by the task's input + and output schema (see below), plus `patient_id`. +- Loaders: `from pyhealth.datasets import get_dataloader`; + `get_dataloader(ds, batch_size=64, shuffle=True)` yields dict batches. +- Models (`from pyhealth.models import ...`): RNN, Transformer, RETAIN, AdaCare, + ConCare, StageNet, TCN, MLP, CNN, LogisticRegression, ... All take + `dataset=train_dataset` as the first argument (it carries the fitted feature + processors and label schema) plus model hyperparameters, e.g. + `RNN(dataset=train_dataset, embedding_dim=128, hidden_dim=128)`. + `model(**batch)` returns a dict with "loss", "y_prob", "y_true", "logit". +- Training: `from pyhealth.trainer import Trainer`; + `trainer = Trainer(model=model, device=device, metrics=[...], + output_path="trainer_output")`; `trainer.train(train_dataloader=..., + val_dataloader=..., epochs=..., monitor=..., optimizer_params={"lr": 1e-3}, + weight_decay=..., patience=...)`. With a val loader and `monitor`, the best + checkpoint is reloaded at the end, but ONLY if logging is enabled (the + default; checkpoints go to `output_path`, inside your working directory). + `enable_logging=False` silently keeps the last epoch instead. Binary metrics include "roc_auc", "pr_auc", "f1"; + multiclass metrics include "accuracy", "f1_macro", "cohen_kappa". +- You may also write a custom `torch.nn.Module` (e.g. subclass + `pyhealth.models.BaseModel`) as long as `model(**batch)` returns a dict with + "y_prob" (and ideally "y_true") for a batch from `get_dataloader`. +""" + +REFINE_INSTRUCTIONS = ( + "Improve your solution. Consider a stronger or different architecture" + " (e.g. Transformer, RETAIN, StageNet, AdaCare) or tuning its capacity.", + "Improve generalization: tune dropout, weight decay, learning rate and the" + " number of epochs, and use early stopping on the validation set.", + "Address class imbalance (e.g. a weighted loss or sampling) if the label" + " distribution is skewed, and make sure the monitored metric matches the" + " scoring metrics.", + "Use the execution output to find inefficiencies or bugs. Make sure" + " training fits comfortably within the time budget.", + "Try combining features differently or ensembling a few small models," + " keeping the total runtime within the time budget.", +) + + +# A working custom-model solution: shows how to wrap a model in +# pyhealth.models.BaseModel and train it with a plain PyTorch loop. Students +# are shown it as a template; tests run it end to end. +BASEMODEL_TEMPLATE = """\ +import copy + +import numpy as np +import torch +import torch.nn as nn +from pyhealth.datasets import get_dataloader +from pyhealth.models import BaseModel +from sklearn.metrics import average_precision_score, roc_auc_score + + +class BagOfCodesModel(BaseModel): + \"\"\"Embeds each token-sequence feature, mean-pools it and classifies.\"\"\" + + def __init__(self, dataset, embedding_dim=128, hidden_dim=128, dropout=0.3): + super().__init__(dataset) # Sets feature_keys, label_keys, mode. + self.label_key = self.label_keys[0] + self.embeddings = nn.ModuleDict() + for key in self.feature_keys: + processor = dataset.input_processors[key] + # Token features: integer ids in [0, processor.size()), 0 = padding. + self.embeddings[key] = nn.Embedding( + processor.size(), embedding_dim, padding_idx=0 + ) + self.head = nn.Sequential( + nn.Linear(len(self.feature_keys) * embedding_dim, hidden_dim), + nn.ReLU(), + nn.Dropout(dropout), + nn.Linear(hidden_dim, self.get_output_size()), + ) + + def forward(self, **batch): + pooled = [] + for key in self.feature_keys: + x = batch[key].to(self.device).long() + x = x.reshape(x.shape[0], -1) # Flatten nested sequences: (B, L). + mask = (x != 0).unsqueeze(-1).float() + emb = self.embeddings[key](x) * mask + pooled.append(emb.sum(1) / mask.sum(1).clamp(min=1.0)) + logits = self.head(torch.cat(pooled, dim=-1)) + y_true = batch[self.label_key].to(self.device) + return { + "loss": self.get_loss_function()(logits, y_true), + "y_prob": self.prepare_y_prob(logits), + "y_true": y_true, + "logit": logits, + } + + +def _evaluate(model, loader): + model.eval() + probs, labels = [], [] + with torch.no_grad(): + for batch in loader: + out = model(**batch) + probs.append(out["y_prob"].cpu().numpy()) + labels.append(out["y_true"].cpu().numpy()) + y_prob, y_true = np.concatenate(probs), np.concatenate(labels) + if model.mode == "binary": + y_prob, y_true = y_prob.ravel(), y_true.ravel() + return {"auroc": roc_auc_score(y_true, y_prob), + "auprc": average_precision_score(y_true, y_prob)} + return {"accuracy": float((y_prob.argmax(1) == y_true.ravel()).mean())} + + +def solution(train_dataset, val_dataset, device, epochs=5): + train_loader = get_dataloader(train_dataset, batch_size=128, shuffle=True) + val_loader = get_dataloader(val_dataset, batch_size=512, shuffle=False) + model = BagOfCodesModel(train_dataset).to(device) + optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4) + best_score, best_state = -np.inf, None + for epoch in range(epochs): + model.train() + for batch in train_loader: + optimizer.zero_grad() + model(**batch)["loss"].backward() + optimizer.step() + metrics = _evaluate(model, val_loader) + score = float(np.mean(list(metrics.values()))) + print(f"epoch {epoch}: {metrics}") + if score > best_score: + best_score, best_state = score, copy.deepcopy(model.state_dict()) + model.load_state_dict(best_state) + return model +""" + + +def _monitor_metric(spec: spec_lib.EnvSpec) -> str: + return "roc_auc" if spec.task_type == spec_lib.BINARY else "accuracy" + + +def _feature_table(prepared: prepare_lib.PreparedEnv) -> str: + """Markdown table of each batch field: processor, vocab and tensor layout.""" + inputs = prepared.schema_summary.get("inputs", {}) + outputs = prepared.schema_summary.get("outputs", {}) + rows = [ + "| batch key | role | processor | vocab/classes | batch value (B=4) |", + "|---|---|---|---|---|", + ] + for key, info in prepared.batch_summary.items(): + meta = inputs.get(key) or outputs.get(key) or {} + role = ( + "input feature" + if key in inputs + else "label" if key in outputs else "metadata" + ) + if "shape" in info: + layout = f"{info['type']} {info['dtype']} {info['shape']}" + elif "shapes" in info: + layout = f"{info['type']} {info['shapes']} {info['dtypes']}" + else: + layout = info["type"] + token = " (token ids, 0 = padding)" if meta.get("is_token") else "" + rows.append( + f"| `{key}` | {role} | {meta.get('processor', '-')}{token} |" + f" {meta.get('size', '-')} | {layout} |" + ) + return "\n".join(rows) + + +def build_task_prompt( + spec: spec_lib.EnvSpec, prepared: prepare_lib.PreparedEnv +) -> str: + """The initial student prompt describing the environment and contract.""" + sizes = prepared.split_sizes + labels = prepared.label_distribution + total = sum(labels.values()) or 1 + label_lines = ", ".join( + f"{k}: {v / total:.1%}" for k, v in sorted(labels.items()) + ) + y_prob_shape = ( + "shape (B, 1) or (B,): probability of the positive class" + if spec.task_type == spec_lib.BINARY + else "shape (B, C): class probabilities" + ) + return f"""\ +# Task: {spec.name} +{spec.description} + +Task type: {spec.task_type}. Label field: "{spec.label_key}". +Label distribution (training split sample): {label_lines}. + +## What to write +``` +def solution(train_dataset, val_dataset, device): + ... + return model +``` +- `train_dataset`, `val_dataset`: PyHealth `SampleDataset` subsets split by + patient ({sizes[prepare_lib.LEARNING_TRAIN]} / {sizes[prepare_lib.LEARNING_VAL]} samples during learning cycles; the final + evaluation reruns your code on the full splits, {sizes[prepare_lib.TRAIN]} / {sizes[prepare_lib.VAL]} samples). +- `device`: "cuda" (one GPU, shared, keep memory use modest) or "cpu". +- Return the trained model. It is evaluated with `model.eval()` and + `torch.no_grad()` on batches from + `get_dataloader(eval_dataset, batch_size={spec.inference_batch_size}, shuffle=False)`: + `model(**batch)` must return a dict with "y_prob" ({y_prob_shape}) and + ideally "y_true". Learning cycles are scored on the validation split; the + final score uses a held-out test split you never see. +- Budget: {int(spec.learning_timeout_seconds)}s per learning-cycle run, {int(spec.final_timeout_seconds)}s for the final run + (training + evaluation), about {spec.epochs_hint} epochs. Runs over budget are killed. +- Print short progress lines (e.g. validation metrics per epoch); they are + shown back to you. + +## Data +Datasets yield dict samples; `get_dataloader(dataset, batch_size, shuffle)` +collates them into dict batches with these fields: + +{_feature_table(prepared)} + +Token sequences are padded with 0 within a batch; vocab sizes include the +padding and unknown tokens, so `nn.Embedding(processor.size(), dim, +padding_idx=0)` is valid. Fitted processors are available as +`train_dataset.input_processors[key]` / `train_dataset.output_processors[key]`. + +## Scoring +{scoring.get_metrics_description(spec.metric_weights)} + +{PYHEALTH_CHEATSHEET} +## Template A: built-in PyHealth model + Trainer +{_strip_plan(build_seed_solution(spec))} +## Template B: custom model wrapped in `pyhealth.models.BaseModel` +Subclass `BaseModel` (call `super().__init__(dataset)`; it provides +`feature_keys`, `label_keys`, `mode`, `device`, `get_output_size()`, +`get_loss_function()` and `prepare_y_prob(logits)`) and return the standard +output dict, so the model works with both a custom loop and `Trainer`: + +```python +{BASEMODEL_TEMPLATE}``` + +These templates are starting points: change the architecture, features, +loss, optimization and training schedule to maximize the score. +""" + + +def _strip_plan(solution_text: str) -> str: + """Returns only the fenced code block of a seed solution.""" + return solution_text[solution_text.index("```python"):] + + +def build_seed_solution(spec: spec_lib.EnvSpec) -> str: + """A working RNN baseline in the solution contract, used as a seed.""" + metrics = ( + '["roc_auc", "pr_auc", "f1"]' + if spec.task_type == spec_lib.BINARY + else '["accuracy", "f1_macro"]' + ) + code = textwrap.dedent(f"""\ + import torch + from pyhealth.datasets import get_dataloader + from pyhealth.models import RNN + from pyhealth.trainer import Trainer + + + def solution(train_dataset, val_dataset, device): + train_loader = get_dataloader(train_dataset, batch_size=64, shuffle=True) + val_loader = get_dataloader(val_dataset, batch_size=256, shuffle=False) + model = RNN( + dataset=train_dataset, + embedding_dim=128, + hidden_dim=128, + num_layers=1, + dropout=0.3, + ) + # Logging must stay enabled for the best checkpoint to be reloaded. + trainer = Trainer( + model=model, + device=device, + metrics={metrics}, + output_path="trainer_output", + ) + trainer.train( + train_dataloader=train_loader, + val_dataloader=val_loader, + epochs={spec.epochs_hint}, + monitor="{_monitor_metric(spec)}", + optimizer_params={{"lr": 1e-3}}, + ) + print("val:", trainer.evaluate(val_loader)) + return model + """) + return ( + "Plan: start from a GRU baseline (PyHealth RNN) trained with the" + " PyHealth Trainer, monitoring the validation metric.\n\n" + f"```python\n{code}```\n" + ) diff --git a/codegen_classroom/pyhealth_env/pyhealth_env_test.py b/codegen_classroom/pyhealth_env/pyhealth_env_test.py new file mode 100644 index 0000000..bb22ab7 --- /dev/null +++ b/codegen_classroom/pyhealth_env/pyhealth_env_test.py @@ -0,0 +1,295 @@ +"""Tests for PyHealth environments using a small synthetic SampleDataset.""" + +import os +from unittest import mock + +from absl.testing import absltest +from codegen_classroom import agents +from codegen_classroom import code_execution +from codegen_classroom import environments +from codegen_classroom import scoring +from codegen_classroom import student_learning +from codegen_classroom.pyhealth_env import environment as environment_lib +from codegen_classroom.pyhealth_env import prepare as prepare_lib +from codegen_classroom.pyhealth_env import prompts +from codegen_classroom.pyhealth_env import registry +from codegen_classroom.pyhealth_env import runner +from codegen_classroom.pyhealth_env import spec as spec_lib +import numpy as np +from pyhealth.datasets import create_sample_dataset + +_NUM_PATIENTS = 60 + + +def _make_samples(): + rng = np.random.default_rng(0) + samples = [] + for p in range(_NUM_PATIENTS): + for v in range(2): + label = int(rng.random() < 0.4) + conditions = [f"cond-{rng.integers(0, 20)}" for _ in range(3)] + if label: + conditions.append("cond-risk") + samples.append({ + "patient_id": f"patient-{p}", + "visit_id": f"visit-{p}-{v}", + "conditions": conditions, + "procedures": [f"proc-{rng.integers(0, 10)}"], + "label": label, + }) + return samples + + +_BAD_RETURN_CODE = """ +def solution(train_dataset, val_dataset, device): + return 42 +""" + +_ERROR_CODE = """ +def solution(train_dataset, val_dataset, device): + raise RuntimeError("boom from student") +""" + + +class PyHealthEnvTest(absltest.TestCase): + + @classmethod + def setUpClass(cls): + super().setUpClass() + cls.sample_dataset = create_sample_dataset( + samples=_make_samples(), + input_schema={"conditions": "sequence", "procedures": "sequence"}, + output_schema={"label": "binary"}, + dataset_name="synthetic", + in_memory=False, + ) + + def setUp(self): + super().setUp() + # Run student code on CPU for determinism and to leave GPUs alone. + self.enter_context( + mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": ""}) + ) + self.spec = spec_lib.EnvSpec( + name="synthetic", + label_key="label", + sample_dataset_path=str(self.sample_dataset.path), + description="Synthetic binary task.", + split_ratios=(0.6, 0.2, 0.2), + max_learning_train_samples=40, + epochs_hint=2, + learning_timeout_seconds=300, + ) + self.env_dir = self.create_tempdir().full_path + self.prepared = prepare_lib.prepare_env(self.spec, self.env_dir) + self.env = environment_lib.PyHealthEnvironment( + self.spec, self.prepared, work_dir=self.create_tempdir().full_path + ) + + def test_split_by_patient_indices(self): + p2i = {f"p{i}": [2 * i, 2 * i + 1] for i in range(10)} + train, val, test = prepare_lib.split_by_patient_indices( + p2i, (0.6, 0.2, 0.2), seed=1 + ) + all_idx = np.concatenate([train, val, test]) + self.assertCountEqual(all_idx.tolist(), list(range(20))) + self.assertLen(train, 12) + # Both samples of a patient land in the same split. + for split in (train, val, test): + self.assertTrue(all((i ^ 1) in set(split) for i in split)) + again = prepare_lib.split_by_patient_indices(p2i, (0.6, 0.2, 0.2), seed=1) + np.testing.assert_array_equal(train, again[0]) + + def test_prepare_env_contents_and_cache(self): + sizes = self.prepared.split_sizes + self.assertEqual( + sizes["train"] + sizes["val"] + sizes["test"], 2 * _NUM_PATIENTS + ) + self.assertEqual(sizes["learning_train"], 40) + self.assertIn("conditions", self.prepared.schema_summary["inputs"]) + self.assertEqual(set(self.prepared.label_distribution), {"0", "1"}) + with mock.patch.object( + prepare_lib, "split_by_patient_indices" + ) as mock_split: + cached = prepare_lib.prepare_env(self.spec, self.env_dir) + mock_split.assert_not_called() + self.assertEqual(cached.split_sizes, sizes) + + def test_execute_seed_solution_learning_phase(self): + code, _ = code_execution.extract_python_code( + prompts.build_seed_solution(self.spec), "solution" + ) + execution = self.env.execute( + code, phase=environments.LEARNING, label="student_1_cycle_0" + ) + n_val = self.prepared.split_sizes["learning_val"] + self.assertEqual(execution.y_pred.shape, (n_val, 1)) + self.assertTrue(execution.y_true.index.equals(execution.y_pred.index)) + self.assertTrue(((execution.y_pred >= 0) & (execution.y_pred <= 1)).all().all()) + self.assertIn("[runner]", execution.execution_feedback) + score, metrics, _, _ = scoring.score_predictions( + execution.y_true, + execution.y_pred, + {}, + self.spec.metric_weights, + binary_threshold=self.spec.metric_binary_threshold, + ) + self.assertTrue(np.isfinite(score)) + self.assertIn("auroc", metrics) + + def test_execute_basemodel_template(self): + execution = self.env.execute( + prompts.BASEMODEL_TEMPLATE, + phase=environments.LEARNING, + label="basemodel_template", + ) + self.assertLen(execution.y_pred, self.prepared.split_sizes["learning_val"]) + self.assertIn("epoch 0", execution.execution_feedback) + + def test_task_prompt_documents_batch_and_templates(self): + prompt = prompts.build_task_prompt(self.spec, self.prepared) + self.assertIn("| `conditions` | input feature | SequenceProcessor", prompt) + self.assertIn("| `label` | label |", prompt) + self.assertIn("class BagOfCodesModel(BaseModel)", prompt) + self.assertIn("from pyhealth.models import RNN", prompt) + self.assertEqual(prompt.count("def solution("), 3) # Contract + 2 templates. + + def test_execute_final_phase_uses_test_split(self): + code, _ = code_execution.extract_python_code( + prompts.build_seed_solution(self.spec), "solution" + ) + execution = self.env.execute( + code, phase=environments.FINAL, label="student_1_best_final" + ) + self.assertLen(execution.y_pred, self.prepared.split_sizes["test"]) + + def test_relative_dirs_are_resolved(self): + cwd = os.getcwd() + base = self.create_tempdir().full_path + os.chdir(base) + try: + prepared = prepare_lib.prepare_env(self.spec, "envs/rel") + env = environment_lib.PyHealthEnvironment( + self.spec, prepared, work_dir="runs/rel" + ) + self.assertTrue(os.path.isabs(prepared.splits_path)) + execution = env.execute( + prompts.BASEMODEL_TEMPLATE, phase=environments.LEARNING, label="rel" + ) + finally: + os.chdir(cwd) + self.assertLen(execution.y_pred, prepared.split_sizes["learning_val"]) + + def test_execute_student_error_raises_execution_error(self): + with self.assertRaisesRegex( + code_execution.ExecutionError, "boom from student" + ): + self.env.execute(_ERROR_CODE, phase=environments.LEARNING, label="err") + + def test_execute_bad_return_raises_execution_error(self): + with self.assertRaisesRegex(code_execution.ExecutionError, "trained model"): + self.env.execute( + _BAD_RETURN_CODE, phase=environments.LEARNING, label="bad" + ) + + def test_learning_cycle_with_env(self): + mock_backend = mock.Mock() + mock_backend.retryable_exceptions = () + mock_backend.fatal_exceptions = () + mock_backend.generate.return_value = "Try again." + with mock.patch.object( + agents.backends, "make_backend", return_value=mock_backend + ): + student = agents.StudentAgent( + agents.AgentConfig(name="student_1"), + self.create_tempdir().full_path, + seed_solution=prompts.build_seed_solution(self.spec), + ) + results = student_learning.run_parallel_classroom_learning_cycle( + current_classroom=[student], + x_train=None, + y_train=None, + x_val=None, + y_val=None, + current_cycle=0, + student_initial_prompt=prompts.build_task_prompt( + self.spec, self.prepared + ), + student_function_name="solution", + metric_weights=self.spec.metric_weights, + early_stopping_patience=2, + student_refine_instruction_options=list(prompts.REFINE_INSTRUCTIONS), + max_mutation_count=1, + mutation_strategy_options=[agents.MutationStrategy.SELF_BEST_CODE], + code_execution_max_retries=0, + metric_binary_threshold=self.spec.metric_binary_threshold, + env=self.env, + ) + self.assertLen(results, 1) + self.assertEqual(results[0].exclude_reason, "") + self.assertTrue(np.isfinite(results[0].overall_score)) + self.assertIn("auroc", results[0].metrics) + self.assertEqual(student.best_cycle, 0) + + +class RunnerScoresTest(absltest.TestCase): + + def test_binary_two_column_probs(self): + y_true, y_score = runner.to_scores( + np.array([[0.0], [1.0]]), + np.array([[0.8, 0.2], [0.1, 0.9]]), + "binary", + ) + np.testing.assert_allclose(y_true, [0, 1]) + np.testing.assert_allclose(y_score, [0.2, 0.9]) + + def test_binary_rejects_non_probabilities(self): + with self.assertRaisesRegex(ValueError, r"\[0, 1\]"): + runner.to_scores(np.array([0, 1]), np.array([[-1.0], [2.0]]), "binary") + + def test_multiclass_argmax(self): + y_true, y_score = runner.to_scores( + np.array([2, 0]), + np.array([[0.1, 0.2, 0.7], [0.6, 0.3, 0.1]]), + "multiclass", + ) + np.testing.assert_allclose(y_true, [2, 0]) + np.testing.assert_allclose(y_score, [2, 0]) + + def test_length_mismatch(self): + with self.assertRaisesRegex(ValueError, "rows"): + runner.to_scores(np.array([0, 1, 1]), np.array([0.5, 0.5]), "binary") + + +class RegistryTest(absltest.TestCase): + + def test_get_spec_merges_dataset_kwargs(self): + spec = registry.get_spec( + "mimic4_mortality", + {"dataset_kwargs": {"dev": True}, "epochs_hint": 3}, + ) + self.assertEqual(spec.name, "mimic4_mortality") + self.assertTrue(spec.dataset_kwargs["dev"]) + self.assertIn("ehr_root", spec.dataset_kwargs) + self.assertEqual(spec.epochs_hint, 3) + self.assertEqual(spec.metric_weights, {"auroc": 0.5, "auprc": 0.5}) + + def test_all_registered_specs_build(self): + for name in registry.list_envs(): + self.assertEqual(registry.get_spec(name).name, name) + + def test_unknown_env(self): + with self.assertRaises(KeyError): + registry.get_spec("nope") + + def test_spec_validation(self): + with self.assertRaises(ValueError): + spec_lib.EnvSpec(name="x", label_key="y") + with self.assertRaises(ValueError): + spec_lib.EnvSpec( + name="x", label_key="y", sample_dataset_path="/p", task_type="bad" + ) + + +if __name__ == "__main__": + absltest.main() diff --git a/codegen_classroom/pyhealth_env/registry.py b/codegen_classroom/pyhealth_env/registry.py new file mode 100644 index 0000000..2443f36 --- /dev/null +++ b/codegen_classroom/pyhealth_env/registry.py @@ -0,0 +1,139 @@ +"""Built-in PyHealth environments mirroring `benchmark_tasks/`. + +Dataset roots and cache directories default to the paths used by the +benchmark scripts on this cluster; override any field from the run config +(e.g. `env.dataset_kwargs.ehr_root`). +""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +import copy +import dataclasses +from typing import Any + +from codegen_classroom.pyhealth_env import spec as spec_lib + +_MIMIC4_ROOT = "/shared/rsaas/physionet.org/files/mimiciv/2.2" +_MIMIC4_CACHE = "/shared/eng/pyhealth_agent/cache" +_MIMIC3_ROOT = "/srv/local/data/MIMIC-III" +_MIMIC3_CACHE = "/shared/eng/pyhealth/mimic3_dataset" +_EICU_ROOT = "/srv/local/data/physionet.org/files/eicu-crd/2.0" +_EICU_CACHE = "/shared/eng/pyhealth/eicu_dataset" + + +def _mimic4(tables, task_class, label_key, task_type, description): + return spec_lib.EnvSpec( + name="", + dataset_class="pyhealth.datasets.MIMIC4Dataset", + dataset_kwargs={ + "ehr_root": _MIMIC4_ROOT, + "ehr_tables": tables, + "cache_dir": _MIMIC4_CACHE, + "num_workers": 4, + }, + task_class=task_class, + label_key=label_key, + task_type=task_type, + description=description, + max_learning_train_samples=20000, + ) + + +_REGISTRY: Mapping[str, Callable[[], spec_lib.EnvSpec]] = { + "mimic4_mortality": lambda: _mimic4( + [ + "patients", + "admissions", + "diagnoses_icd", + "procedures_icd", + "prescriptions", + ], + "pyhealth.tasks.MortalityPredictionMIMIC4", + "mortality", + spec_lib.BINARY, + "In-hospital mortality prediction from the next admission on" + " MIMIC-IV (diagnoses, procedures, prescriptions).", + ), + "mimic4_readmission": lambda: _mimic4( + [ + "patients", + "admissions", + "diagnoses_icd", + "procedures_icd", + "prescriptions", + ], + "pyhealth.tasks.ReadmissionPredictionMIMIC4", + "readmission", + spec_lib.BINARY, + "Hospital readmission prediction on MIMIC-IV.", + ), + "mimic4_los": lambda: _mimic4( + ["admissions", "diagnoses_icd", "procedures_icd", "prescriptions"], + "pyhealth.tasks.LengthOfStayPredictionMIMIC4", + "los", + spec_lib.MULTICLASS, + "Length-of-stay category prediction (multiclass) on MIMIC-IV.", + ), + "mimic3_mortality": lambda: spec_lib.EnvSpec( + name="", + dataset_class="pyhealth.datasets.MIMIC3Dataset", + dataset_kwargs={ + "root": _MIMIC3_ROOT, + "tables": ["DIAGNOSES_ICD", "PROCEDURES_ICD", "PRESCRIPTIONS"], + "cache_dir": _MIMIC3_CACHE, + }, + task_class="pyhealth.tasks.MortalityPredictionMIMIC3", + label_key="mortality", + task_type=spec_lib.BINARY, + description="Mortality prediction from the next visit on MIMIC-III.", + ), + "eicu_mortality": lambda: spec_lib.EnvSpec( + name="", + dataset_class="pyhealth.datasets.eICUDataset", + dataset_kwargs={ + "root": _EICU_ROOT, + "tables": ["diagnosis", "medication", "physicalexam"], + "cache_dir": _EICU_CACHE, + }, + task_class="pyhealth.tasks.MortalityPredictionEICU", + label_key="mortality", + task_type=spec_lib.BINARY, + description="ICU mortality prediction on eICU.", + ), +} + + +def list_envs() -> list[str]: + return sorted(_REGISTRY) + + +def get_spec( + name: str, overrides: Mapping[str, Any] | None = None +) -> spec_lib.EnvSpec: + """Returns the named built-in spec, with optional field overrides. + + `dataset_kwargs` and `task_kwargs` overrides are merged into the defaults + rather than replacing them. + + Args: + name: Registered environment name (see `list_envs()`), or "custom" to build + the spec entirely from `overrides`. + overrides: EnvSpec field overrides. + + Returns: + The EnvSpec. + """ + overrides = dict(overrides or {}) + if name == "custom": + return spec_lib.EnvSpec(name=overrides.pop("name", "custom"), **overrides) + if name not in _REGISTRY: + raise KeyError(f"Unknown PyHealth env {name!r}. Available: {list_envs()}") + spec = _REGISTRY[name]() + spec.name = name + for key in ("dataset_kwargs", "task_kwargs"): + if key in overrides: + merged = copy.deepcopy(dict(getattr(spec, key))) + merged.update(overrides.pop(key) or {}) + overrides[key] = merged + return dataclasses.replace(spec, **overrides) diff --git a/codegen_classroom/pyhealth_env/runner.py b/codegen_classroom/pyhealth_env/runner.py new file mode 100644 index 0000000..f168a74 --- /dev/null +++ b/codegen_classroom/pyhealth_env/runner.py @@ -0,0 +1,199 @@ +"""Child-process entrypoint that trains and evaluates one student solution. + +Invoked by `PyHealthEnvironment` as + + python -m codegen_classroom.pyhealth_env.runner --job job.json + +in a fresh interpreter (never forked), with CUDA_VISIBLE_DEVICES pinned to the +GPU slot assigned to this execution. The runner: + + 1. Reloads the prepared SampleDataset and the train/val/eval index splits. + 2. Imports the student's code and calls + `model = solution(train_dataset, val_dataset, device)`. + 3. Runs inference itself on the evaluation split, which the student never + sees, collecting `y_prob` and `y_true` from the same batches so they are + aligned regardless of dataloader ordering. + 4. Writes `y_true` and `y_score` to the job's output .npz file. + +Any exception propagates to stderr and a non-zero exit code, which the parent +turns into an ExecutionError fed back to the student. +""" + +from __future__ import annotations + +import argparse +import importlib.util +import json +import os +import random +import sys +import time +from typing import Any + +import numpy as np + + +def _set_seed(seed: int) -> None: + import torch # pylint: disable=g-import-not-at-top + + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + + +def _load_solution(code_file: str, function_name: str): + module_spec = importlib.util.spec_from_file_location( + "student_solution", code_file + ) + module = importlib.util.module_from_spec(module_spec) + sys.modules["student_solution"] = module + module_spec.loader.exec_module(module) + fn = getattr(module, function_name, None) + if fn is None or not callable(fn): + raise AttributeError( + f"Your code must define a callable `{function_name}(train_dataset," + " val_dataset, device)`." + ) + return fn + + +def _to_numpy(x: Any) -> np.ndarray: + if hasattr(x, "detach"): + x = x.detach().float().cpu().numpy() + return np.asarray(x) + + +def run_inference( + model: Any, + eval_dataset: Any, + label_key: str, + batch_size: int, +) -> tuple[np.ndarray, np.ndarray]: + """Returns (y_true, y_prob) over `eval_dataset` using `model(**batch)`.""" + import torch # pylint: disable=g-import-not-at-top + from pyhealth.datasets import get_dataloader # pylint: disable=g-import-not-at-top + + if not callable(model): + raise TypeError( + "solution() must return a trained model callable as `model(**batch)`" + f" (e.g. a pyhealth.models.BaseModel); got {type(model).__name__}." + ) + if isinstance(model, torch.nn.Module): + model.eval() + loader = get_dataloader(eval_dataset, batch_size=batch_size, shuffle=False) + y_true_parts, y_prob_parts = [], [] + with torch.no_grad(): + for batch in loader: + output = model(**batch) + if not isinstance(output, dict) or "y_prob" not in output: + raise TypeError( + "model(**batch) must return a dict containing 'y_prob' (as" + " pyhealth.models.BaseModel does); got" + f" {type(output).__name__} with keys" + f" {list(output) if isinstance(output, dict) else None}." + ) + y_true = output.get("y_true", batch[label_key]) + y_prob_parts.append(_to_numpy(output["y_prob"])) + y_true_parts.append(_to_numpy(y_true)) + return np.concatenate(y_true_parts), np.concatenate(y_prob_parts) + + +def to_scores( + y_true: np.ndarray, y_prob: np.ndarray, task_type: str +) -> tuple[np.ndarray, np.ndarray]: + """Converts raw model outputs to 1-D (y_true, y_score) for scoring. + + Binary: y_score is the positive-class probability. Multiclass: y_score is the + predicted class index (argmax of y_prob). + + Args: + y_true: Labels as returned by the model or batch. + y_prob: Probabilities as returned by the model. + task_type: "binary" or "multiclass". + + Returns: + Tuple of 1-D (y_true, y_score) arrays of equal length. + """ + n = len(y_true) + if len(y_prob) != n: + raise ValueError( + f"y_prob has {len(y_prob)} rows but there are {n} labels." + ) + if task_type == "binary": + if y_prob.ndim == 2 and y_prob.shape[1] == 2: + y_prob = y_prob[:, 1] + y_score = y_prob.reshape(n, -1) + if y_score.shape[1] != 1: + raise ValueError( + "Binary task: y_prob must have shape (n,), (n, 1) or (n, 2); got" + f" {y_prob.shape}." + ) + y_score = y_score[:, 0] + y_true = y_true.reshape(n, -1)[:, 0] + if np.any((y_score < 0) | (y_score > 1)): + raise ValueError("Binary task: y_prob must be probabilities in [0, 1].") + else: + if y_prob.ndim != 2: + raise ValueError( + f"Multiclass task: y_prob must have shape (n, C); got {y_prob.shape}." + ) + y_score = y_prob.argmax(axis=1) + y_true = ( + y_true.argmax(axis=1) + if y_true.ndim == 2 and y_true.shape[1] > 1 + else y_true.reshape(-1) + ) + if not np.all(np.isfinite(y_score)): + raise ValueError("Predictions contain NaN or infinite values.") + return y_true.astype(np.float64), y_score.astype(np.float64) + + +def main(argv=None) -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--job", required=True) + args = parser.parse_args(argv) + with open(args.job) as f: + job = json.load(f) + + import torch # pylint: disable=g-import-not-at-top + from pyhealth.datasets import SampleDataset # pylint: disable=g-import-not-at-top + + # SampleDataset.close() deletes its directory, which is the shared cache. + SampleDataset.close = lambda self: None + + _set_seed(job.get("seed", 0)) + splits = np.load(job["splits_path"]) + dataset = SampleDataset(path=job["sample_dataset_path"]) + train_dataset = dataset.subset(splits[job["train_split"]].tolist()) + val_dataset = dataset.subset(splits[job["val_split"]].tolist()) + eval_dataset = dataset.subset(splits[job["eval_split"]].tolist()) + device = "cuda" if torch.cuda.is_available() else "cpu" + print( + f"[runner] train={len(train_dataset)} val={len(val_dataset)}" + f" device={device}" + f" CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES')}", + flush=True, + ) + + solution = _load_solution(job["code_file"], job["function_name"]) + t0 = time.time() + model = solution(train_dataset, val_dataset, device) + train_time = time.time() - t0 + + t0 = time.time() + y_true, y_prob = run_inference( + model, eval_dataset, job["label_key"], job["inference_batch_size"] + ) + y_true, y_score = to_scores(y_true, y_prob, job["task_type"]) + print( + f"[runner] solution() took {train_time:.1f}s; inference on" + f" {len(y_score)} samples took {time.time() - t0:.1f}s.", + flush=True, + ) + np.savez(job["output_file"], y_true=y_true, y_score=y_score) + + +if __name__ == "__main__": + main() diff --git a/codegen_classroom/pyhealth_env/spec.py b/codegen_classroom/pyhealth_env/spec.py new file mode 100644 index 0000000..32cd06d --- /dev/null +++ b/codegen_classroom/pyhealth_env/spec.py @@ -0,0 +1,88 @@ +"""Specification of a PyHealth training environment (dataset + task).""" + +from __future__ import annotations + +from collections.abc import Mapping +import dataclasses +import importlib +from typing import Any + +BINARY = "binary" +MULTICLASS = "multiclass" +TASK_TYPES = (BINARY, MULTICLASS) + + +@dataclasses.dataclass +class EnvSpec: + """A PyHealth dataset + task that students write model code for. + + Students write `solution(train_dataset, val_dataset, device)` returning a + trained model; the environment runs inference on the evaluation split and + scores the probabilities. + """ + + name: str + # Label field in the task's output_schema, e.g. "mortality". + label_key: str + task_type: str = BINARY + # Fully qualified classes, e.g. "pyhealth.datasets.MIMIC4Dataset". + dataset_class: str | None = None + dataset_kwargs: Mapping[str, Any] = dataclasses.field(default_factory=dict) + task_class: str | None = None + task_kwargs: Mapping[str, Any] = dataclasses.field(default_factory=dict) + # Use a pre-built SampleDataset directory (containing schema.pkl) instead of + # building one from dataset_class/task_class. + sample_dataset_path: str | None = None + description: str = "" + # Patient-level train/val/test split. + split_ratios: tuple[float, float, float] = (0.8, 0.1, 0.1) + split_seed: int = 42 + # Caps on samples used during learning cycles (None = use all). Final + # evaluation always trains on the full train split and scores the full test + # split. + max_learning_train_samples: int | None = None + max_learning_val_samples: int | None = None + # Wall-clock limits per execution, in seconds. + learning_timeout_seconds: float = 1800 + final_timeout_seconds: float = 7200 + # Suggested epoch budget mentioned in the student prompt. + epochs_hint: int = 5 + # Weighted metrics used for the overall score. + metric_weights: Mapping[str, float] | None = None + # Batch size used by the runner for inference on the evaluation split. + inference_batch_size: int = 256 + + def __post_init__(self): + if self.task_type not in TASK_TYPES: + raise ValueError( + f"task_type must be one of {TASK_TYPES}, got {self.task_type!r}." + ) + if self.sample_dataset_path is None and not ( + self.dataset_class and self.task_class + ): + raise ValueError( + "EnvSpec needs either sample_dataset_path or both dataset_class and" + " task_class." + ) + self.split_ratios = tuple(self.split_ratios) + if self.metric_weights is None: + self.metric_weights = default_metric_weights(self.task_type) + + @property + def metric_binary_threshold(self) -> float | None: + """Threshold for binarizing probabilities in binary-label metrics.""" + return 0.5 if self.task_type == BINARY else None + + +def default_metric_weights(task_type: str) -> dict[str, float]: + if task_type == BINARY: + return {"auroc": 0.5, "auprc": 0.5} + return {"accuracy": 0.5, "f1_macro": 0.5} + + +def import_object(qualified_name: str) -> Any: + """Imports `pkg.module.Name` and returns `Name`.""" + module_name, _, attr = qualified_name.rpartition(".") + if not module_name: + raise ValueError(f"Expected a fully qualified name, got {qualified_name!r}") + return getattr(importlib.import_module(module_name), attr) diff --git a/codegen_classroom/reporting.py b/codegen_classroom/reporting.py index a1c86ef..5b668b5 100644 --- a/codegen_classroom/reporting.py +++ b/codegen_classroom/reporting.py @@ -23,7 +23,6 @@ import logging import os import re -import sys from typing import Any import warnings @@ -152,122 +151,36 @@ def suppress_common_warnings() -> None: ) -def is_colab() -> bool: - """Returns True if the code is running in a Google Colab notebook.""" - return "google.colab" in sys.modules - - -def is_notebook() -> bool: - """Returns True if running in a Jupyter-like notebook environment (including Colab).""" - if is_colab(): - return True - try: - from IPython import get_ipython # pylint: disable=g-import-not-at-top - - shell_obj = get_ipython() - if not shell_obj: - return False - shell = shell_obj.__class__.__name__ - if shell == "ZMQInteractiveShell": - return True - if shell == "TerminalInteractiveShell": - return False - except (NameError, AttributeError, ImportError): - pass - return False - - def render_markdown(text: str) -> None: - """Renders Markdown text in notebooks, or prints plain text in terminals.""" - if is_notebook(): - from IPython import display as ipython_display # pylint: disable=g-import-not-at-top - - ipython_display.display(ipython_display.Markdown(text)) - else: - print(text) + """Prints Markdown text to the console.""" + print(text) def render_code(code: str, language: str = "python") -> None: """Renders a fenced code block.""" - fenced = f"```{language}\n{code}\n```" - if is_notebook(): - from IPython import display as ipython_display # pylint: disable=g-import-not-at-top - - ipython_display.display(ipython_display.Markdown(fenced)) - else: - print(fenced) + print(f"```{language}\n{code}\n```") -def colorized_html_diff(text1: str, text2: str, dark_mode: bool = False) -> str: - """For Colab: generates a colorized HTML diff between two code texts.""" - diff = difflib.ndiff( - text1.splitlines(keepends=True), text2.splitlines(keepends=True) +def render_code_diff(old: str, new: str) -> None: + """Prints a unified diff between two code texts.""" + diff = difflib.unified_diff( + old.splitlines(keepends=False), + new.splitlines(keepends=False), + lineterm="", ) - result = [] - add_color = "#223922" if dark_mode else "#d4fcdc" - del_color = "#451e1e" if dark_mode else "#ffe0e0" - - for line in diff: - sign = line[0] - content = line[2:] - escaped = ( - content.replace("&", "&").replace("<", "<").replace(">", ">") - ) - if sign == " ": - result.append(f" {escaped}") - elif sign == "+": - result.append( - f"+ {escaped}" - ) - elif sign == "-": - result.append( - f"- {escaped}" - ) - return "".join(result) - - -def render_code_diff(old: str, new: str, dark_mode: bool = False) -> None: - """Renders a code diff (colored HTML in notebooks, unified diff text in CLI).""" - if is_notebook(): - from IPython import display as ipython_display # pylint: disable=g-import-not-at-top - - code_diff_html = colorized_html_diff(old, new, dark_mode=dark_mode) - ipython_display.display( - ipython_display.HTML( - f"
{code_diff_html}
" - ) - ) - else: - diff = difflib.unified_diff( - old.splitlines(keepends=False), - new.splitlines(keepends=False), - lineterm="", - ) - print("\n".join(diff)) + print("\n".join(diff)) def render_json(data: dict[str, Any]) -> None: """Renders JSON data.""" json_str = json.dumps(data, indent=2) - if is_notebook(): - from IPython import display as ipython_display # pylint: disable=g-import-not-at-top - - ipython_display.display( - ipython_display.Markdown(f"```json\n{json_str}\n```") - ) - else: - print(f"```json\n{json_str}\n```") + print(f"```json\n{json_str}\n```") def render_dataframe(df: pd.DataFrame, precision: int | None = 3) -> None: """Renders a pandas DataFrame.""" df_display = df.round(precision) if precision is not None else df - if is_notebook(): - from IPython import display as ipython_display # pylint: disable=g-import-not-at-top - - ipython_display.display(df_display) - else: - print(df_display.to_string()) + print(df_display.to_string()) def render_leaderboard(path: str | epath.Path) -> None: @@ -937,16 +850,13 @@ def build_student_report( def render_report( sections: list[ReportSection], - dark_mode: bool = False, precision: int = 3, ) -> None: - """Renders a report to the notebook or terminal.""" + """Prints a report to the console (plots are only written by write_report).""" for sec in sections: if sec.title: heading = "#" * sec.level render_markdown(f"{heading} {sec.title}") - if sec.kind != "plot": - pass if sec.kind == "markdown": if sec.payload: render_markdown(str(sec.payload)) @@ -956,42 +866,14 @@ def render_report( render_code(str(sec.payload), language="text") elif sec.kind == "code_diff": old, new = sec.payload - render_code_diff(old, new, dark_mode=dark_mode) + render_code_diff(old, new) elif sec.kind == "dataframe": render_dataframe(sec.payload, precision=precision) elif sec.kind == "plot": - # In notebook, plot live. Headless, note it skip. - if is_notebook(): - if sec.payload["type"] == "progress": - plotting.plot_progress( - sec.payload["data"], - metric_names=["overall_score"], - backend="plotly", - dark_mode=dark_mode, - show=True, - ) - elif sec.payload["type"] == "final": - plotting.plot_final_test( - metrics=["overall_score"], - final_results_path=sec.payload["data"], - backend="plotly", - dark_mode=dark_mode, - show=True, - ) - else: - print("(Plot skipped in terminal display; noted as saved)") + print("(Plot saved to the run directory by write_report.)") elif sec.kind == "keyvalue": - items = { - k: v for k, v in sec.payload.items() if not str(k).startswith("_") - } - if is_notebook(): - bullets = "\n".join([ - f"* **{k}:** {_format_cell_value(v, precision=precision)}" - for k, v in items.items() - ]) - render_markdown(bullets) - else: - for k, v in items.items(): + for k, v in sec.payload.items(): + if not str(k).startswith("_"): print(f"{k}: {_format_cell_value(v, precision=precision)}") diff --git a/codegen_classroom/run.py b/codegen_classroom/run.py new file mode 100644 index 0000000..c8a3f20 --- /dev/null +++ b/codegen_classroom/run.py @@ -0,0 +1,425 @@ +"""Headless end-to-end classroom run driven by a YAML config. + +Usage (see scripts/run_classroom.py and configs/): + + python scripts/run_classroom.py --config configs/mimic4_mortality_vllm.yaml \ + --override run.num_cycles=2 task.overrides.dataset_kwargs.dev=true + +Steps: set up the run directory and logging, prepare the task, build the +students, run the learning cycles, run the final evaluation, analyze the best +solution, and write `report.md` (+ plots) into the run directory. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +import copy +import dataclasses +import logging +import os +from typing import Any + +from codegen_classroom import agent_config +from codegen_classroom import agents +from codegen_classroom import data_utils +from codegen_classroom import environments +from codegen_classroom import reporting +from codegen_classroom import scoring +from codegen_classroom import student_learning +import numpy as np +import pandas as pd + +logger = logging.getLogger("codegen_classroom.run") + +DEFAULT_CONFIG: dict[str, Any] = { + "run": { + "output_dir": "runs", + "name": None, # Defaults to the task name. + "seed": 0, + "num_cycles": 5, + "final_strategies": ["first", "last", "best"], + "leaderboard_path": None, + "analyze_best": True, + }, + "llm": { + "backend": agent_config.DEFAULT_BACKEND, + "base_url": None, + "api_key": None, + # One student per entry: a model name, or {backend, model, base_url}. + "models": None, + "temperature": 0.8, + "max_output_tokens": 16384, + "extra_body": None, + "max_code_generation_retries": None, + "retry_sleep_seconds": None, + "analysis_model": None, # Defaults to the first student's model. + }, + "classroom": { + "early_stopping_patience": 2, + "max_mutation_count": 2, + "mutation_strategies": ["self_best_code", "best_peer_crossover"], + "code_execution_max_retries": 2, + "classroom_feedback": True, + # Seed with the baseline solution: "none", "first" or "all" students. + "seed_solution": "first", + "num_workers": None, + }, + "task": { + "type": "pyhealth", # "pyhealth" or "toy_regression". + # pyhealth: + "env": None, # Registered env name or "custom". + "overrides": {}, # EnvSpec field overrides. + "env_cache_dir": None, # Defaults to /envs/. + "gpus": [], + "jobs_per_gpu": 1, + # toy_regression: + "n_samples": 300, + "code_execution_timeout_seconds": 60, + }, +} + +_TOY_SYSTEM_INSTRUCTION = ( + "You are an expert ML engineer. Answer with a short plan followed by one" + " complete ```python code block defining the requested function." +) +_TOY_PROMPT = """\ +Write a Python function `solution(x_train, y_train, x_val)` that trains a +regression model on (x_train, y_train) and returns a tuple +`(y_pred, feedback)`, where `y_pred` is a pandas DataFrame of predictions for +x_val with the same index as x_val and the single column 'target', and +`feedback` is a short string describing what you did. pandas is available as +`pd`, numpy as `np`; scikit-learn may be imported. +""" +_TOY_REFINE = ( + "Improve your solution, e.g. with feature engineering or regularization.", + "Try a different model family and compare.", +) + + +def deep_merge(base: Mapping[str, Any], update: Mapping[str, Any]) -> dict: + """Recursively merges `update` into a copy of `base`.""" + merged = copy.deepcopy(dict(base)) + for key, value in (update or {}).items(): + if isinstance(value, Mapping) and isinstance(merged.get(key), Mapping): + merged[key] = deep_merge(merged[key], value) + else: + merged[key] = copy.deepcopy(value) + return merged + + +def apply_overrides(config: dict, overrides: Sequence[str]) -> dict: + """Applies `a.b.c=value` overrides; values are parsed as YAML.""" + import yaml # pylint: disable=g-import-not-at-top + + config = copy.deepcopy(config) + for item in overrides or []: + if "=" not in item: + raise ValueError(f"Override must look like key.path=value, got {item!r}") + path, raw = item.split("=", 1) + keys = path.strip().split(".") + node = config + for key in keys[:-1]: + if not isinstance(node.get(key), dict): + node[key] = {} # Missing or None-valued defaults become mappings. + node = node[key] + node[keys[-1]] = yaml.safe_load(raw) + return config + + +def load_config( + path: str | None, overrides: Sequence[str] | None = None +) -> dict: + import yaml # pylint: disable=g-import-not-at-top + + user = {} + if path: + with open(path) as f: + user = yaml.safe_load(f) or {} + return apply_overrides(deep_merge(DEFAULT_CONFIG, user), overrides or []) + + +@dataclasses.dataclass +class TaskSetup: + """Everything the learning loop needs for one task.""" + + name: str + system_instruction: str + initial_prompt: str + refine_instructions: Sequence[str] + metric_weights: Mapping[str, float] + metric_binary_threshold: float | None = None + seed_solution: str | None = None + env: environments.TaskEnvironment | None = None + timeout_seconds: float | None = None + # DataFrame tasks only. + data: Mapping[str, pd.DataFrame | None] = dataclasses.field( + default_factory=dict + ) + + +def _setup_pyhealth_task(config: dict, run_dir: str) -> TaskSetup: + from codegen_classroom import pyhealth_env # pylint: disable=g-import-not-at-top + + task_cfg = config["task"] + if not task_cfg.get("env"): + raise ValueError("task.env must name a PyHealth env (or 'custom').") + spec = pyhealth_env.registry.get_spec( + task_cfg["env"], task_cfg.get("overrides") + ) + env_dir = task_cfg.get("env_cache_dir") or os.path.join( + config["run"]["output_dir"], "envs", spec.name + ) + env, prepared = pyhealth_env.make_environment( + spec, + env_dir=env_dir, + work_dir=run_dir, + gpus=task_cfg.get("gpus") or [], + jobs_per_gpu=task_cfg.get("jobs_per_gpu", 1), + seed=config["run"]["seed"], + ) + return TaskSetup( + name=spec.name, + system_instruction=pyhealth_env.prompts.SYSTEM_INSTRUCTION, + initial_prompt=pyhealth_env.prompts.build_task_prompt(spec, prepared), + refine_instructions=pyhealth_env.prompts.REFINE_INSTRUCTIONS, + metric_weights=spec.metric_weights, + metric_binary_threshold=spec.metric_binary_threshold, + seed_solution=pyhealth_env.prompts.build_seed_solution(spec), + env=env, + ) + + +def _setup_toy_regression_task(config: dict) -> TaskSetup: + rng = np.random.default_rng(config["run"]["seed"]) + n = int(config["task"]["n_samples"]) + x = pd.DataFrame(rng.normal(size=(n, 3)), columns=["x1", "x2", "x3"]) + y = pd.DataFrame({ + "target": 2 * x["x1"] - x["x2"] ** 2 + 0.1 * rng.normal(size=n) + }) + i_val, i_test = int(0.6 * n), int(0.8 * n) + return TaskSetup( + name="toy_regression", + system_instruction=_TOY_SYSTEM_INSTRUCTION, + initial_prompt=_TOY_PROMPT, + refine_instructions=_TOY_REFINE, + metric_weights={"r2": 1.0}, + timeout_seconds=config["task"]["code_execution_timeout_seconds"], + data={ + "x_train": x.iloc[:i_val], + "y_train": y.iloc[:i_val], + "x_val": x.iloc[i_val:i_test], + "y_val": y.iloc[i_val:i_test], + "x_test": x.iloc[i_test:], + "y_test": y.iloc[i_test:], + }, + ) + + +def setup_task(config: dict, run_dir: str) -> TaskSetup: + task_type = config["task"]["type"] + if task_type == "pyhealth": + return _setup_pyhealth_task(config, run_dir) + if task_type == "toy_regression": + return _setup_toy_regression_task(config) + raise ValueError(f"Unknown task.type {task_type!r}.") + + +def build_student_configs( + config: dict, system_instruction: str +) -> list[agents.AgentConfig]: + """One AgentConfig per entry of llm.models (plain names or dicts).""" + llm = config["llm"] + models = llm.get("models") + if not models: + raise ValueError("llm.models must list at least one model.") + configs = [] + for entry in models: + if isinstance(entry, str): + entry = {"model": entry} + cfg = student_learning.get_classroom_student_configs( + system_instruction=system_instruction, + api_key=entry.get("api_key", llm.get("api_key")), + model_names=[entry["model"]], + temperature=entry.get("temperature", llm["temperature"]), + max_output_tokens=entry.get( + "max_output_tokens", llm["max_output_tokens"] + ), + max_code_generation_retries=llm.get("max_code_generation_retries"), + retry_sleep_seconds=llm.get("retry_sleep_seconds"), + backend=entry.get("backend", llm["backend"]), + base_url=entry.get("base_url", llm.get("base_url")), + extra_body=entry.get("extra_body", llm.get("extra_body")), + )[0] + configs.append(cfg) + # Names are student_1_; renumber so they stay unique. + for i, cfg in enumerate(configs): + cfg.name = cfg.name.replace("student_1_", f"student_{i + 1}_", 1) + return configs + + +def _analyze_best( + config: dict, + task: TaskSetup, + run_dir: str, + students: Mapping[str, agents.StudentAgent], + eval_results: Sequence[Mapping[str, Any]], + student_configs: Sequence[agents.AgentConfig], +) -> None: + best_name, best_data = student_learning.get_best_solution_from_eval_results( + eval_results + ) + if not best_name or best_name not in students: + logger.warning("No best solution to analyze.") + return + student = students[best_name] + best_result = student.cycle_results[student.best_cycle] + base = student_configs[0] + analysis_cfg = dataclasses.replace( + base, + name="solution_analyst", + system_instruction="", + model_name=config["llm"].get("analysis_model") or base.model_name, + temperature=0.2, + response_mime_type=None, + ) + analyst = agents.SolutionAnalysisAgent(analysis_cfg, run_base_dir=run_dir) + analysis = analyst.analyze( + student_name=best_name, + best_code=best_result.code, + best_plan=best_result.plan, + best_test_metrics=best_data, + best_val_metrics=best_result.metrics, + student_instructions=task.initial_prompt, + metrics_description=scoring.get_metrics_description(task.metric_weights), + task_description_for_prompt=task.name, + ) + logger.info("Best solution analysis: %s", analysis) + + +def run(config: dict) -> str: + """Runs a full classroom experiment. Returns the run directory.""" + run_cfg, cls_cfg = config["run"], config["classroom"] + run_name = run_cfg.get("name") or config["task"].get("env") or "run" + run_dir = os.path.join( + data_utils.resolve_workspace_dir(run_cfg["output_dir"]), + f"{run_name}_{data_utils.generate_run_id()}", + ) + data_utils.create_dir(run_dir) + reporting.setup_logger( + "codegen_classroom", log_file=os.path.join(run_dir, "run.log") + ) + reporting.suppress_common_warnings() + import yaml # pylint: disable=g-import-not-at-top + + with open(os.path.join(run_dir, "config.yaml"), "w") as f: + yaml.safe_dump(config, f, sort_keys=False) + logger.info("Run directory: %s", run_dir) + + task = setup_task(config, run_dir) + student_configs = build_student_configs(config, task.system_instruction) + seed_mode = cls_cfg.get("seed_solution", "none") + classroom: student_learning.Classroom = [] + for i, cfg in enumerate(student_configs): + seeded = seed_mode == "all" or (seed_mode == "first" and i == 0) + classroom.append( + agents.StudentAgent( + cfg, + run_dir, + seed_solution=task.seed_solution if seeded else None, + ) + ) + students = {s.name: s for s in classroom} + history = {s.name: [] for s in classroom} + mutation_strategies = [ + agents.MutationStrategy[name.upper()] + for name in cls_cfg["mutation_strategies"] + ] + metric_keys = ["overall_score", *task.metric_weights] + + active = list(classroom) + for cycle in range(int(run_cfg["num_cycles"])): + if not active: + logger.info("All students excluded; stopping after %d cycles.", cycle) + break + feedback = ( + scoring.get_classroom_feedback_string( + cycle, active, history, metric_keys + ) + if cls_cfg.get("classroom_feedback") + else None + ) + results = student_learning.run_parallel_classroom_learning_cycle( + current_classroom=active, + x_train=task.data.get("x_train"), + y_train=task.data.get("y_train"), + x_val=task.data.get("x_val"), + y_val=task.data.get("y_val"), + current_cycle=cycle, + student_initial_prompt=task.initial_prompt, + student_function_name="solution", + metric_weights=task.metric_weights, + early_stopping_patience=cls_cfg["early_stopping_patience"], + student_refine_instruction_options=list(task.refine_instructions), + max_mutation_count=cls_cfg["max_mutation_count"], + mutation_strategy_options=mutation_strategies, + code_execution_max_retries=cls_cfg["code_execution_max_retries"], + code_execution_timeout_seconds=task.timeout_seconds, + classroom_feedback=feedback or None, + num_workers=cls_cfg.get("num_workers"), + metric_binary_threshold=task.metric_binary_threshold, + env=task.env, + ) + active = student_learning.process_cycle_results( + active, results, history, cycle + ) + + final_results_path = os.path.join(run_dir, "final_eval_results.json") + eval_results, _ = student_learning.run_parallel_classroom_final_evaluation( + students=students, + x_learning=task.data.get("x_train"), + y_learning=task.data.get("y_train"), + x_test_final=task.data.get("x_test"), + y_test_final=task.data.get("y_test"), + student_function_name="solution", + metric_weights=task.metric_weights, + metric_binary_threshold=task.metric_binary_threshold, + code_execution_timeout_seconds=task.timeout_seconds, + enforce_timeout=task.env is None, + strategies=run_cfg["final_strategies"], + num_workers=cls_cfg.get("num_workers"), + final_results_path=final_results_path, + predictions_dir=os.path.join(run_dir, "predictions"), + env=task.env, + ) + + if run_cfg.get("analyze_best"): + try: + _analyze_best( + config, task, run_dir, students, eval_results, student_configs + ) + except Exception as e: # pylint: disable=broad-exception-caught + logger.warning("Solution analysis failed: %s", e) + + sections = reporting.build_run_report( + run_dir, students=classroom, metric_weights=task.metric_weights + ) + reporting.write_report(run_dir, sections, metric_weights=task.metric_weights) + reporting.render_report(sections) + + if run_cfg.get("leaderboard_path"): + best_name, best_data = student_learning.get_best_solution_from_eval_results( + eval_results + ) + if best_name: + row = { + "run_id": os.path.basename(run_dir), + "task": task.name, + "student_name": best_name, + "model_name": students[best_name].model_name, + **{k: v for k, v in (best_data or {}).items() if k != "traceback"}, + } + data_utils.submit_results_to_local_leaderboard( + run_cfg["leaderboard_path"], pd.DataFrame([row]) + ) + logger.info("Done. Report: %s", os.path.join(run_dir, "report.md")) + return run_dir diff --git a/codegen_classroom/run_test.py b/codegen_classroom/run_test.py new file mode 100644 index 0000000..df1d9dc --- /dev/null +++ b/codegen_classroom/run_test.py @@ -0,0 +1,105 @@ +"""Tests for the headless run driver.""" + +import json +import os +from unittest import mock + +from absl.testing import absltest +from codegen_classroom import agents +from codegen_classroom import run + +_RIDGE_SOLUTION = """Plan: ridge regression. + +```python +from sklearn.linear_model import Ridge + +def solution(x_train, y_train, x_val): + model = Ridge().fit(x_train, y_train) + y_pred = pd.DataFrame(model.predict(x_val), index=x_val.index, columns=y_train.columns) + return y_pred, "Trained Ridge" +``` +""" + + +class ConfigTest(absltest.TestCase): + + def test_load_config_merges_and_overrides(self): + path = self.create_tempfile( + content='run:\n num_cycles: 7\nllm:\n models: [a]\n' + ).full_path + config = run.load_config( + path, ["run.num_cycles=2", "task.gpus=[1, 2]", "llm.extra_body.top_p=0.9"] + ) + self.assertEqual(config["run"]["num_cycles"], 2) + self.assertEqual(config["run"]["seed"], 0) # Default kept. + self.assertEqual(config["task"]["gpus"], [1, 2]) + self.assertEqual(config["llm"]["extra_body"], {"top_p": 0.9}) + self.assertEqual(config["llm"]["models"], ["a"]) + + def test_bad_override(self): + with self.assertRaises(ValueError): + run.apply_overrides({}, ["no_equals_sign"]) + + def test_build_student_configs_mixed_backends(self): + config = run.load_config(None) + config["llm"].update( + backend="vllm", + base_url="http://h:1/v1", + models=["Qwen/Coder", "Qwen/Coder", {"backend": "gemini", "model": "g"}], + ) + cfgs = run.build_student_configs(config, "sys") + self.assertEqual( + [c.name for c in cfgs], + ["student_1_Qwen_Coder", "student_2_Qwen_Coder", "student_3_g"], + ) + self.assertEqual([c.backend for c in cfgs], ["vllm", "vllm", "gemini"]) + self.assertEqual(cfgs[0].base_url, "http://h:1/v1") + self.assertEqual(cfgs[0].system_instruction, "sys") + + +class ToyRunTest(absltest.TestCase): + + def test_end_to_end_toy_regression(self): + mock_backend = mock.Mock() + mock_backend.retryable_exceptions = () + mock_backend.fatal_exceptions = () + mock_backend.generate.side_effect = lambda history, prompt, **kw: ( + json.dumps({"model_tags": "ridge"}) + if kw.get("json_mode") + else _RIDGE_SOLUTION + ) + output_dir = self.create_tempdir().full_path + config = run.load_config( + None, + [ + f"run.output_dir={output_dir}", + "run.num_cycles=2", + "llm.models=[m1, m2]", + "llm.backend=vllm", + "task.type=toy_regression", + "classroom.seed_solution=none", + ], + ) + with mock.patch.object( + agents.backends, "make_backend", return_value=mock_backend + ): + run_dir = run.run(config) + + for name in ( + "config.yaml", + "run.log", + "final_eval_results.json", + "best_analysis.json", + "report.md", + ): + self.assertTrue( + os.path.exists(os.path.join(run_dir, name)), f"missing {name}" + ) + with open(os.path.join(run_dir, "final_eval_results.json")) as f: + results = json.load(f) + self.assertLen(results, 2) + self.assertGreater(results[0]["best"]["overall_score"], 0.0) + + +if __name__ == "__main__": + absltest.main() diff --git a/codegen_classroom/student_learning.py b/codegen_classroom/student_learning.py index 9942c84..71d45c3 100644 --- a/codegen_classroom/student_learning.py +++ b/codegen_classroom/student_learning.py @@ -22,6 +22,7 @@ import logging import os import random +import re import time import traceback from typing import Any @@ -30,6 +31,7 @@ from codegen_classroom import agents from codegen_classroom import code_execution from codegen_classroom import data_utils +from codegen_classroom import environments from codegen_classroom import scoring import numpy as np import pandas as pd @@ -362,6 +364,7 @@ def run_learning_cycle( metric_binary_threshold: float | None = None, verbose_logging: bool = False, render_feedback_as_md: bool = False, + env: environments.TaskEnvironment | None = None, ) -> agents.StudentCycleResult: """Runs a single learning cycle for a given student. @@ -401,6 +404,9 @@ def run_learning_cycle( regression predictions to binary labels for classification metrics. verbose_logging: If True, log full code and other verbose details. render_feedback_as_md: If True, render execution feedback as markdown. + env: Optional TaskEnvironment that executes the code instead of the + DataFrame contract. When set, x/y data, auxiliary data, cv_splits and + validator are ignored. Returns: A StudentCycleResult object containing the results of the cycle. @@ -476,7 +482,20 @@ def run_learning_cycle( ) # Execute the extracted code over splits or single set - if cv_splits is not None: + if env is not None: + execution = env.execute( + extracted_code, + phase=environments.LEARNING, + label=f"{student.name}_cycle_{current_cycle}_attempt_{attempt}", + timeout_seconds=code_execution_timeout_seconds, + ) + y_pred_df = execution.y_pred + execution_feedback = execution.execution_feedback + exec_time = execution.exec_time + scoring_y_val = execution.y_true + scoring_aux_val = execution.auxiliary + + elif cv_splits is not None: logger.debug( "[%s] Executing code with %d CV splits (OOF)...", student.name, @@ -677,6 +696,7 @@ def run_parallel_classroom_learning_cycle( metric_binary_threshold: float | None = None, verbose_logging: bool = False, render_feedback_as_md: bool = False, + env: environments.TaskEnvironment | None = None, ) -> ClassroomResults: """Runs the learning cycle for all students in the classroom in parallel. @@ -711,23 +731,26 @@ def run_parallel_classroom_learning_cycle( regression predictions to binary labels for classification metrics. verbose_logging: If True, log full code and other verbose details. render_feedback_as_md: If True, render execution feedback as markdown. + env: Optional TaskEnvironment that executes student code (see + `run_learning_cycle`). DataFrame inputs may be None when set. Returns: A sequence of StudentCycleResult objects, one for each student. """ if not current_classroom: return [] - dataframe_inputs = { - "x_train": x_train, - "y_train": y_train, - "x_val": x_val, - "y_val": y_val, - } - if auxiliary_train_data is not None: - dataframe_inputs.update(auxiliary_train_data) - if auxiliary_val_data is not None: - dataframe_inputs.update(auxiliary_val_data) - _validate_dataframe_inputs(**dataframe_inputs) + if env is None: + dataframe_inputs = { + "x_train": x_train, + "y_train": y_train, + "x_val": x_val, + "y_val": y_val, + } + if auxiliary_train_data is not None: + dataframe_inputs.update(auxiliary_train_data) + if auxiliary_val_data is not None: + dataframe_inputs.update(auxiliary_val_data) + _validate_dataframe_inputs(**dataframe_inputs) parallel_run_args = [ { "student": student, @@ -757,6 +780,7 @@ def run_parallel_classroom_learning_cycle( "metric_binary_threshold": metric_binary_threshold, "verbose_logging": verbose_logging, "render_feedback_as_md": render_feedback_as_md, + "env": env, } for student in current_classroom ] @@ -845,6 +869,7 @@ def run_final_evaluation( auxiliary_test_data: Mapping[str, pd.DataFrame] | None = None, validator: scoring.ValidatorFunction = scoring.execution_validator, cv_splits: Sequence[tuple[np.ndarray, np.ndarray]] | None = None, + env: environments.TaskEnvironment | None = None, ) -> tuple[Mapping[str, Any], Mapping[str, pd.DataFrame]]: """Runs the final evaluation for a single student on the fixed test set. @@ -879,6 +904,10 @@ def run_final_evaluation( validator: Function to validate the output of the executed code. cv_splits: Optional sequence of tuples of indices (train_idx, val_idx) to evaluate OOF predictions in CV fashion. + env: Optional TaskEnvironment that executes the code on its held-out + split. When set, x/y data, auxiliary data, cv_splits and validator are + ignored, and the environment's own final timeout applies unless + enforce_timeout is set. Returns: A tuple containing: @@ -963,7 +992,19 @@ def run_final_evaluation( ) # Execute the extracted code. - if cv_splits is not None: + if env is not None: + execution = env.execute( + extracted_code, + phase=environments.FINAL, + label=f"{student_name}_{strategy}_final_eval", + timeout_seconds=effective_timeout, + ) + y_pred_df = execution.y_pred + exec_time = execution.exec_time + y_test_scored = execution.y_true + aux_test_scored = execution.auxiliary + + elif cv_splits is not None: y_pred_df, _, exec_time = _execute_with_cv_splits( extracted_code=extracted_code, target_fn_name=student_function_name, @@ -1030,6 +1071,7 @@ def run_final_evaluation( except ( code_execution.ExtractionError, code_execution.ExecutionError, + code_execution.ExecutionTimeoutError, scoring.ValidatorError, scoring.ScoreError, scoring.EvaluatorError, @@ -1097,6 +1139,7 @@ def run_parallel_classroom_final_evaluation( final_results_path: str | None = None, predictions_dir: str | None = None, cv_splits: Sequence[tuple[np.ndarray, np.ndarray]] | None = None, + env: environments.TaskEnvironment | None = None, ) -> tuple[Sequence[Mapping[str, Any]], Mapping[str, pd.DataFrame]]: """Runs final evaluation for all students in parallel and saves results. @@ -1128,6 +1171,8 @@ def run_parallel_classroom_final_evaluation( None, predictions are not saved. cv_splits: Optional sequence of tuples of indices (train_idx, val_idx) to evaluate OOF predictions in CV fashion. + env: Optional TaskEnvironment that executes student code (see + `run_final_evaluation`). DataFrame inputs may be None when set. Returns: A tuple containing: @@ -1142,17 +1187,18 @@ def run_parallel_classroom_final_evaluation( if strategies is None: strategies = ["best"] - dataframe_inputs = { - "x_learning": x_learning, - "y_learning": y_learning, - "x_test_final": x_test_final, - "y_test_final": y_test_final, - } - if auxiliary_learning_data is not None: - dataframe_inputs.update(auxiliary_learning_data) - if auxiliary_test_data is not None: - dataframe_inputs.update(auxiliary_test_data) - _validate_dataframe_inputs(**dataframe_inputs) + if env is None: + dataframe_inputs = { + "x_learning": x_learning, + "y_learning": y_learning, + "x_test_final": x_test_final, + "y_test_final": y_test_final, + } + if auxiliary_learning_data is not None: + dataframe_inputs.update(auxiliary_learning_data) + if auxiliary_test_data is not None: + dataframe_inputs.update(auxiliary_test_data) + _validate_dataframe_inputs(**dataframe_inputs) logger.info( "Starting Final Evaluation for %d students (Strategies: %s)", @@ -1178,6 +1224,7 @@ def run_parallel_classroom_final_evaluation( "validator": validator, "cv_splits": cv_splits, "metric_binary_threshold": metric_binary_threshold, + "env": env, }) num_workers = num_workers if num_workers is not None else len(students) @@ -1202,7 +1249,6 @@ def run_parallel_classroom_final_evaluation( # Save predictions if directory is provided if predictions_dir and all_predictions: - prediction_col = y_test_final.columns[0] data_utils.create_dir(predictions_dir) logger.info( "Saving %d Predictions to CSV in %s", @@ -1214,8 +1260,9 @@ def run_parallel_classroom_final_evaluation( if y_pred_df is None or y_pred_df.empty: logger.warning("Skipping empty/None prediction for '%s'", pred_key) continue - y_pred_df = y_pred_df.set_index(x_test_final.index) - y_pred_df.columns = [prediction_col] + if env is None: + y_pred_df = y_pred_df.set_index(x_test_final.index) + y_pred_df.columns = [y_test_final.columns[0]] pred_path = os.path.join(predictions_dir, f"{pred_key}.csv") data_utils.save_dataframe_to_csv(y_pred_df, pred_path, index=True) logger.debug("Saved prediction for '%s' to: %s", pred_key, pred_path) @@ -1319,33 +1366,49 @@ def get_best_solution_from_eval_results( def get_classroom_student_configs( system_instruction: str, api_key: str | None = None, - model_names: Sequence[str] | None = None, + model_names: Sequence[str | tuple[str, str]] | None = None, temperature: float = 1.0, max_output_tokens: int = 32768, max_code_generation_retries: int | None = None, retry_sleep_seconds: int | None = None, + backend: str | None = None, + base_url: str | None = None, + extra_body: Mapping[str, Any] | None = None, ) -> Sequence[agents.AgentConfig]: """Returns a sequence of student configurations for a classroom. Args: system_instruction: The base system instruction to use for each agent. api_key: The API key for the generative model service. - model_names: Optional sequence of model names to use. + model_names: Optional sequence of model names to use. Entries may also be + `(backend, model_name)` tuples to mix backends in one classroom (e.g. + `[("vllm", "Qwen/Qwen3-Coder-30B-A3B-Instruct-FP8"), ("gemini", + "gemini-3.5-flash")]`). temperature: The default temperature for the models. max_output_tokens: The default maximum output tokens for the models. max_code_generation_retries: Optional maximum number of retries for code generation. retry_sleep_seconds: Optional sleep time in seconds between retries. + backend: Backend for plain-string `model_names` entries ("gemini" or + "openai_compat"/"vllm"). Defaults to `agent_config.DEFAULT_BACKEND`. + base_url: Server URL for OpenAI-compatible backends. + extra_body: Extra request fields for OpenAI-compatible backends. Returns: A sequence of agents.AgentConfig objects. """ if model_names is None: model_names = agent_config.DEFAULT_STUDENT_MODELS + default_backend = backend or agent_config.DEFAULT_BACKEND student_cfgs = [] - for i, model_name in enumerate(model_names): - clean_model_name = model_name.replace("-", "_") + for i, entry in enumerate(model_names): + if isinstance(entry, str): + student_backend, model_name = default_backend, entry + else: + student_backend, model_name = entry + # Student names are used in file paths, so keep them filesystem-safe. + clean_model_name = re.sub(r"[^A-Za-z0-9_.]", "_", model_name) student_name = f"student_{i+1}_{clean_model_name}" kwargs = { "name": student_name, @@ -1354,6 +1417,9 @@ def get_classroom_student_configs( "max_output_tokens": max_output_tokens, "system_instruction": system_instruction, "model_name": model_name, + "backend": student_backend, + "base_url": base_url, + "extra_body": extra_body, } if max_code_generation_retries is not None: kwargs["max_code_generation_retries"] = max_code_generation_retries diff --git a/codegen_classroom/student_learning_test.py b/codegen_classroom/student_learning_test.py index b1b13b7..2ae7cc3 100644 --- a/codegen_classroom/student_learning_test.py +++ b/codegen_classroom/student_learning_test.py @@ -46,9 +46,9 @@ def setUp(self): super().setUp() self.base_dir = self.create_tempdir().full_path - # Mock genai Client for agents - self.mock_genai_client = self.enter_context( - mock.patch.object(agents, 'genai', autospec=True) + # Mock the LLM backend for agents + self.mock_make_backend = self.enter_context( + mock.patch.object(agents.backends, 'make_backend', autospec=True) ) # Common data diff --git a/configs/mimic4_mortality_cached_qwen.yaml b/configs/mimic4_mortality_cached_qwen.yaml new file mode 100644 index 0000000..8d66a41 --- /dev/null +++ b/configs/mimic4_mortality_cached_qwen.yaml @@ -0,0 +1,55 @@ +# MIMIC-IV mortality on an existing PyHealth task cache, students served by the +# 4-bit Qwen3.8-27B through vLLM. +# +# # vllm env (see scripts/serve_vllm.sh); the bnb-4bit model is ~21 GB on disk +# MODEL=unsloth/Qwen3.8-27B-unsloth-bnb-4bit GPUS= \ +# bash scripts/serve_vllm.sh +# # PyHealthAgent env +# python scripts/run_classroom.py --config configs/mimic4_mortality_cached_qwen.yaml +# +# Validated without the LLM (seed solution + BaseModel template) on 2026-10-16: +# final test AUROC 0.78 / AUPRC 0.088 for the RNN baseline. + +run: + output_dir: runs + num_cycles: 4 + final_strategies: [first, last, best] + leaderboard_path: runs/leaderboard.tsv + +llm: + backend: vllm + base_url: http://localhost:8000/v1 + models: + - unsloth/Qwen3.8-27B-unsloth-bnb-4bit + - unsloth/Qwen3.8-27B-unsloth-bnb-4bit + - unsloth/Qwen3.8-27B-unsloth-bnb-4bit + temperature: 0.7 + max_output_tokens: 8192 + retry_sleep_seconds: 10 + # Qwen3.x thinks by default; blocks are stripped from responses but + # count against max_output_tokens. Disable to save tokens: + # extra_body: {chat_template_kwargs: {enable_thinking: false}} + +classroom: + early_stopping_patience: 2 + max_mutation_count: 1 + mutation_strategies: [self_best_code, best_peer_crossover] + code_execution_max_retries: 2 + seed_solution: first + +task: + type: pyhealth + env: mimic4_mortality + # Student training GPUs; each A6000 here has only ~7-8 GB free next to the + # resident SGLang servers, enough for one small EHR model per GPU. + gpus: [6, 7] + jobs_per_gpu: 1 + env_cache_dir: runs/envs/mimic4_mortality_cached + overrides: + # Reuse the cached MortalityPredictionMIMIC4 SampleDataset (128,236 visits, + # 57,701 patients; inputs: conditions, procedures, drugs). + sample_dataset_path: /shared/eng/pyhealth_agent/cache/b77effea-2e99-5b1b-90c3-0fa6d2540af1/tasks/MortalityPredictionMIMIC4_187a839c-4f67-585a-bbb3-355429e27594/samples_cdbbc602-34e2-5a41-8643-4c76b08829f6.ld + max_learning_train_samples: 20000 + learning_timeout_seconds: 900 + final_timeout_seconds: 3600 + epochs_hint: 5 diff --git a/configs/mimic4_mortality_mixed.yaml b/configs/mimic4_mortality_mixed.yaml new file mode 100644 index 0000000..85fa9d1 --- /dev/null +++ b/configs/mimic4_mortality_mixed.yaml @@ -0,0 +1,20 @@ +# Mixed classroom: local vLLM students plus a Gemini student (needs +# GEMINI_API_KEY in .env and `pip install -e ".[gemini]"`). + +run: + output_dir: /shared/eng/pyhealth_agent/codegen_runs + num_cycles: 5 + +llm: + backend: vllm + base_url: http://localhost:8000/v1 + models: + - Qwen/Qwen3-Coder-30B-A3B-Instruct-FP8 + - Qwen/Qwen3-Coder-30B-A3B-Instruct-FP8 + - {backend: gemini, model: gemini-3.5-flash} + analysis_model: Qwen/Qwen3-Coder-30B-A3B-Instruct-FP8 + +task: + type: pyhealth + env: mimic4_mortality + gpus: [1, 2] diff --git a/configs/mimic4_mortality_qwen9b_2080ti.yaml b/configs/mimic4_mortality_qwen9b_2080ti.yaml new file mode 100644 index 0000000..42bde88 --- /dev/null +++ b/configs/mimic4_mortality_qwen9b_2080ti.yaml @@ -0,0 +1,53 @@ +# MIMIC-IV mortality (cached task) with Qwen3.5-9B AWQ students on RTX 2080 Ti +# (sunlab-serv-02, 8x 11 GB). The LLM uses GPUs 0-1 (TP=2), students train on +# the rest. +# +# PATH=~/miniconda3/envs/vllm/bin:$PATH MODEL=QuantTrio/Qwen3.5-9B-AWQ \ +# GPUS=0,1 TP=2 MAX_MODEL_LEN=32768 \ +# EXTRA_ARGS='--dtype half --max-num-seqs 16 --limit-mm-per-prompt {"image":0,"video":0}' \ +# bash scripts/serve_vllm.sh +# python scripts/run_classroom.py --config configs/mimic4_mortality_qwen9b_2080ti.yaml +# +# Turing notes: no bf16 (--dtype half), and the default max-num-seqs (256) +# exceeds the linear-attention cache blocks, so lower it. Startup takes ~6 min +# (warmup ~5 min). The KV cache holds ~155k tokens, so a 32k context fits; at +# 16k, within-cycle retry history overflowed the context. + +run: + output_dir: runs + num_cycles: 2 + final_strategies: [best] + leaderboard_path: runs/leaderboard.tsv + +llm: + backend: vllm + base_url: http://localhost:8000/v1 + models: + - QuantTrio/Qwen3.5-9B-AWQ + - QuantTrio/Qwen3.5-9B-AWQ + temperature: 0.7 + max_output_tokens: 4096 + retry_sleep_seconds: 10 + # Thinking off: reasoning plus the ~5k-token prompt and within-cycle retry + # history can still overflow the context. + extra_body: {chat_template_kwargs: {enable_thinking: false}} + +classroom: + early_stopping_patience: 2 + max_mutation_count: 1 + mutation_strategies: [self_best_code, best_peer_crossover] + code_execution_max_retries: 2 + seed_solution: none # Make both students write their own first solution. + +task: + type: pyhealth + env: mimic4_mortality + gpus: [2, 3] + jobs_per_gpu: 1 + env_cache_dir: runs/envs/mimic4_mortality_cached + overrides: + sample_dataset_path: /shared/eng/pyhealth_agent/cache/b77effea-2e99-5b1b-90c3-0fa6d2540af1/tasks/MortalityPredictionMIMIC4_187a839c-4f67-585a-bbb3-355429e27594/samples_cdbbc602-34e2-5a41-8643-4c76b08829f6.ld + max_learning_train_samples: 20000 + learning_timeout_seconds: 900 + final_timeout_seconds: 3600 + epochs_hint: 5 diff --git a/configs/mimic4_mortality_vllm.yaml b/configs/mimic4_mortality_vllm.yaml new file mode 100644 index 0000000..0b39f3b --- /dev/null +++ b/configs/mimic4_mortality_vllm.yaml @@ -0,0 +1,45 @@ +# MIMIC-IV in-hospital mortality with students served by a local vLLM server. +# +# bash scripts/serve_vllm.sh # in the `vllm` env, GPU 0 +# python scripts/run_classroom.py --config configs/mimic4_mortality_vllm.yaml +# +# GPU layout (3x A6000): GPU 0 serves the LLM, GPUs 1-2 train student models. + +run: + output_dir: /shared/eng/pyhealth_agent/codegen_runs + num_cycles: 5 + final_strategies: [first, last, best] + leaderboard_path: /shared/eng/pyhealth_agent/codegen_runs/leaderboard.tsv + +llm: + backend: vllm + base_url: http://localhost:8000/v1 + # One student per entry. Repeating a model gives independent students. + models: + - Qwen/Qwen3-Coder-30B-A3B-Instruct-FP8 + - Qwen/Qwen3-Coder-30B-A3B-Instruct-FP8 + - Qwen/Qwen3-Coder-30B-A3B-Instruct-FP8 + - Qwen/Qwen3-Coder-30B-A3B-Instruct-FP8 + temperature: 0.8 + max_output_tokens: 16384 + retry_sleep_seconds: 10 + # extra_body: {top_p: 0.95} + +classroom: + early_stopping_patience: 2 + max_mutation_count: 2 + mutation_strategies: [self_best_code, best_peer_crossover] + code_execution_max_retries: 2 + seed_solution: first # student 1 starts from the PyHealth RNN baseline + +task: + type: pyhealth + env: mimic4_mortality + gpus: [1, 2] + jobs_per_gpu: 1 + overrides: + max_learning_train_samples: 20000 + learning_timeout_seconds: 1800 + final_timeout_seconds: 7200 + epochs_hint: 5 + # dataset_kwargs: {dev: true} # small MIMIC-IV subset for quick runs diff --git a/configs/regression_toy.yaml b/configs/regression_toy.yaml new file mode 100644 index 0000000..7c099a3 --- /dev/null +++ b/configs/regression_toy.yaml @@ -0,0 +1,25 @@ +# Smoke test: one vLLM student, one cycle, toy regression (no PyHealth data). +# +# python scripts/run_classroom.py --config configs/regression_toy.yaml + +run: + output_dir: runs + num_cycles: 1 + final_strategies: [best] + analyze_best: false + +llm: + backend: vllm + base_url: http://localhost:8000/v1 + models: [Qwen/Qwen3-Coder-30B-A3B-Instruct-FP8] + max_output_tokens: 4096 + max_code_generation_retries: 2 + retry_sleep_seconds: 5 + +classroom: + code_execution_max_retries: 1 + seed_solution: none + +task: + type: toy_regression + n_samples: 300 diff --git a/configs/run_classroom.sub b/configs/run_classroom.sub new file mode 100644 index 0000000..5e9191a --- /dev/null +++ b/configs/run_classroom.sub @@ -0,0 +1,22 @@ +# HTCondor submission for a classroom run (mirrors the repo-root htcondor.sub). +# condor_submit configs/run_classroom.sub CONFIG=configs/mimic4_mortality_vllm.yaml +# The vLLM server must be reachable from the job (same machine or llm.base_url). +CONFIG ?= configs/mimic4_mortality_vllm.yaml + +executable = /bin/bash +arguments = "-c 'source ~/miniconda3/etc/profile.d/conda.sh && conda activate PyHealthAgent && python scripts/run_classroom.py --config $(CONFIG)'" +initialdir = /home/johnwu3/projects/PyHealth_Branch_Testing/codegen-classroom-pyhealth +getenv = True + +Requirements = (Machine == "sunlab-c02.cs.illinois.edu") + +output = logs/classroom.$(Cluster).$(Process).out +error = logs/classroom.$(Cluster).$(Process).err +log = logs/classroom.$(Cluster).$(Process).log + +universe = vanilla +request_cpus = 16 +request_memory = 200GB +request_gpus = 2 + +queue diff --git a/pyproject.toml b/pyproject.toml index 3c562f7..d36eb80 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta" [project] name = "codegen-classroom" version = "0.1.0" -description = "A framework for iterative code generation for ML and data analysis tasks using parallel LLM agents in Colab." +description = "A framework for iterative code generation for ML and data analysis tasks using parallel LLM agents." readme = "README.md" requires-python = ">=3.9" license = {file = "LICENSE"} @@ -19,7 +19,6 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", ] dependencies = [ - "google-genai", "httpx", "numpy", "pandas", @@ -34,9 +33,19 @@ dependencies = [ ] [project.optional-dependencies] -notebook = [ - "plotly", - "ipython", +gemini = [ + "google-genai", +] +openai = [ + "openai>=1.40", +] +pyhealth = [ + "pyhealth", + "torch", + "pyarrow", +] +cli = [ + "pyyaml", ] torch = [ "torch", @@ -52,9 +61,44 @@ dev = [ "pytest", "pytest-mock", ] +# Exact versions of the classroom env used for the reported PyHealth results +# (conda env `PyHealthAgent`, Python 3.12). PyHealth itself was an editable +# install of evanfebrianto/PyHealth@c232681 (branch feat/ehrmamba, reports +# 2.0.1); see README "Reproducing our results". Install with: +# pip install -e ".[repro]" +repro = [ + "torch==2.7.1", + "pydantic==2.11.10", # pyhealth pins pydantic~=2.11.7 + "pydantic-core==2.33.2", + "litdata==0.2.59", + "polars==1.35.2", + "pyarrow==22.0.0", + "numpy==2.2.6", + "pandas==2.3.3", + "scikit-learn==1.7.2", + "scipy==1.17.0", + "matplotlib==3.10.8", + "seaborn==0.13.2", + "absl-py==2.5.1", + "etils==1.14.0", + "httpx==0.28.1", + "python-dotenv==1.2.4", + "openai==3.31.0", + "google-genai==2.8.0", # newest release compatible with pydantic 2.11 + "pyyaml==6.0.3", + "pytest==9.1.1", +] +# The vLLM server must live in its OWN environment: vLLM 0.31 needs torch 2.13 +# (CUDA 13), which conflicts with PyHealth's torch 2.7.1. In a separate env: +# pip install "codegen-classroom[vllm-server]" (or the two pins directly) +vllm-server = [ + "vllm==0.31.0", + "bitsandbytes==0.50.2", +] all = [ - "plotly", - "ipython", + "google-genai", + "openai>=1.40", + "pyyaml", "torch", "torchvision", "jax", @@ -64,5 +108,9 @@ all = [ ] +[tool.setuptools.packages.find] +# configs/ and scripts/ are not packages; only ship the library. +include = ["codegen_classroom*"] + [project.urls] Homepage = "https://github.com/google-research/codegen-classroom" diff --git a/scripts/run_classroom.py b/scripts/run_classroom.py new file mode 100755 index 0000000..1ef05e6 --- /dev/null +++ b/scripts/run_classroom.py @@ -0,0 +1,70 @@ +#!/usr/bin/env python +"""Runs a codegen classroom experiment from a YAML config (no notebook needed). + +Examples: + # PyHealth MIMIC-IV mortality with students served by a local vLLM server. + python scripts/run_classroom.py --config configs/mimic4_mortality_vllm.yaml + + # Quick smoke test against the vLLM server on a toy regression task. + python scripts/run_classroom.py --config configs/regression_toy.yaml + + # Override any config value (values are parsed as YAML). + python scripts/run_classroom.py --config configs/mimic4_mortality_vllm.yaml \ + --override run.num_cycles=2 'task.gpus=[1]' \ + task.overrides.dataset_kwargs.dev=true + + # List the built-in PyHealth environments. + python scripts/run_classroom.py --list-envs +""" + +import argparse +import sys + + +def main() -> int: + # Score lines are print()ed; keep them in order with the logs when redirected. + sys.stdout.reconfigure(line_buffering=True) + parser = argparse.ArgumentParser(description=__doc__.split("\n")[0]) + parser.add_argument("--config", help="Path to a YAML run config.") + parser.add_argument( + "--override", + nargs="*", + default=[], + metavar="KEY=VALUE", + help="Dotted config overrides, e.g. run.num_cycles=3.", + ) + parser.add_argument( + "--list-envs", + action="store_true", + help="Print the built-in PyHealth environments and exit.", + ) + parser.add_argument( + "--print-config", + action="store_true", + help="Print the resolved config and exit.", + ) + args = parser.parse_args() + + if args.list_envs: + from codegen_classroom.pyhealth_env import registry # pylint: disable=g-import-not-at-top + + for name in registry.list_envs(): + spec = registry.get_spec(name) + print(f"{name:24s} {spec.task_type:10s} {spec.description}") + return 0 + + from codegen_classroom import run # pylint: disable=g-import-not-at-top + + config = run.load_config(args.config, args.override) + if args.print_config: + import yaml # pylint: disable=g-import-not-at-top + + print(yaml.safe_dump(config, sort_keys=False)) + return 0 + run_dir = run.run(config) + print(f"Run directory: {run_dir}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/serve_vllm.sh b/scripts/serve_vllm.sh new file mode 100755 index 0000000..3ecca7c --- /dev/null +++ b/scripts/serve_vllm.sh @@ -0,0 +1,47 @@ +#!/usr/bin/env bash +# Starts a vLLM OpenAI-compatible server for codegen classroom students. +# +# Run it in its own conda env (vLLM pins its own torch, which conflicts with +# PyHealth), e.g.: +# conda create -n vllm python=3.12 -y && conda activate vllm && pip install vllm +# bash scripts/serve_vllm.sh +# +# Then point the classroom at it (default in configs/*_vllm.yaml): +# llm.backend: vllm, llm.base_url: http://localhost:8000/v1 +# and check it with: curl http://localhost:8000/v1/models +# +# Settings (environment variables): +# MODEL HF model id or local path to serve. +# GPUS CUDA_VISIBLE_DEVICES for the server (keep student GPUs free). +# TP Tensor-parallel size (number of GPUs in GPUS). +# PORT Server port. +# MAX_MODEL_LEN Context length; must exceed prompt + llm.max_output_tokens. +# GPU_MEM_UTIL Fraction of each GPU's memory vLLM may use. +# EXTRA_ARGS Any extra `vllm serve` arguments. +# +# Pre-quantized bitsandbytes checkpoints (e.g. unsloth/*-bnb-4bit) are detected +# from their config. Text-only use of a multimodal checkpoint can skip the +# vision encoder's memory with EXTRA_ARGS='--limit-mm-per-prompt {"image":0,"video":0}'. +# Example: the 4-bit Qwen3.8-27B checkpoint is ~21 GB (vision tower not +# quantized), so it needs roughly 25-30 GB with KV cache, e.g. one free A6000: +# MODEL=unsloth/Qwen3.8-27B-unsloth-bnb-4bit GPUS=0 MAX_MODEL_LEN=32768 \ +# bash scripts/serve_vllm.sh +# GPU_MEM_UTIL is a fraction of TOTAL memory and must fit in what is free. +set -euo pipefail + +MODEL="${MODEL:-Qwen/Qwen3-Coder-30B-A3B-Instruct-FP8}" +GPUS="${GPUS:-0}" +TP="${TP:-1}" +PORT="${PORT:-8000}" +MAX_MODEL_LEN="${MAX_MODEL_LEN:-65536}" +GPU_MEM_UTIL="${GPU_MEM_UTIL:-0.90}" +EXTRA_ARGS="${EXTRA_ARGS:-}" + +echo "Serving ${MODEL} on GPU(s) ${GPUS} (tp=${TP}) at http://localhost:${PORT}/v1" +# shellcheck disable=SC2086 +CUDA_VISIBLE_DEVICES="${GPUS}" exec vllm serve "${MODEL}" \ + --port "${PORT}" \ + --tensor-parallel-size "${TP}" \ + --max-model-len "${MAX_MODEL_LEN}" \ + --gpu-memory-utilization "${GPU_MEM_UTIL}" \ + ${EXTRA_ARGS}