fix(eval): release scored eval samples' generated media to free CPU mem

.github/workflows/_run-ci.yml

@@ -202,8 +202,8 @@ jobs:
       - name: Set up Python
         uses: actions/setup-python@v5
         with:
-          # 3.10 matches miles-D's setup.py / pyproject (python_requires=">=3.10").
-          python-version: '3.10'
+          # 3.12 matches the CI/dev container image and miles-D's setup.py.
+          python-version: '3.12'
 
       - name: Install uv
         uses: astral-sh/setup-uv@v5

.github/workflows/pre-commit.yml

@@ -22,7 +22,7 @@ jobs:
       - name: Set up Python
         uses: actions/setup-python@v5
         with:
-          python-version: '3.10'
+          python-version: '3.12'
           cache: 'pip'
 
       - name: Install pre-commit

miles/rollout/sglang_diffusion_rollout.py

@@ -381,6 +381,15 @@ async def generate_rollout_async(
 EVAL_PROMPT_DATASET = {}
 
 
+def release_eval_sample_media(args: Namespace, sample: Sample) -> None:
+    """Drop a scored eval sample's media, which only image logging and debug dumps still read."""
+    if args.save_debug_rollout_data is not None:
+        return
+    if sample.index < args.diffusion_log_images:
+        return
+    sample.generated_output = None
+
+
 # eval only
 async def eval_rollout(args: Namespace, rollout_id: int) -> tuple[dict[str, dict[str, list[Any]]], list[list[Sample]]]:
     assert not args.group_rm, "Group RM is not supported for eval rollout"
@@ -456,6 +465,8 @@ async def eval_rollout_single_dataset(
                 "eval_rollout_single_dataset example data, prompt: " f"{[str(row.prompt)]} " f"reward={row.reward}"
             )
             do_print = False
+        for row in rows:
+            release_eval_sample_media(args, row)
         data.extend(rows)
         pbar.update(1)
     pbar.close()

pyproject.toml

@@ -10,7 +10,7 @@ build-backend = "setuptools.build_meta"
 profile = "black"  # black-compatible
 line_length = 119  # should match black parameters
 ignore_whitespace = true  # ignore whitespace for compatibility with the initial style
-py_version = 310  # python 3.10 as a target version
+py_version = 312  # python 3.12 as a target version
 sections = ["FUTURE", "STDLIB", "THIRDPARTY", "FIRSTPARTY", "LOCALFOLDER"]
 default_section = "THIRDPARTY"
 extend_skip = ["setup.py", "docs/source/conf.py"]

setup.py

@@ -41,10 +41,8 @@ def get_tag(self):
             "torch>=2.0",
         ]
     },
-    python_requires=">=3.10",
+    python_requires=">=3.12",
     classifiers=[
-        "Programming Language :: Python :: 3.10",
-        "Programming Language :: Python :: 3.11",
         "Programming Language :: Python :: 3.12",
         "Environment :: GPU :: NVIDIA CUDA",
         "Topic :: Scientific/Engineering :: Artificial Intelligence",

tests/fast/rollout/test_release_eval_sample_media.py

@@ -0,0 +1,38 @@
+from tests.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=15, suite="stage-a-cpu", labels=[])
+
+from argparse import Namespace
+
+import torch
+
+from miles.rollout.sglang_diffusion_rollout import release_eval_sample_media
+from miles.utils.types import Sample
+
+
+def make_sample(index: int) -> Sample:
+    return Sample(index=index, generated_output=torch.zeros(3, 1, 4, 4))
+
+
+def test_media_released_when_not_logged():
+    sample = make_sample(index=3)
+
+    release_eval_sample_media(Namespace(save_debug_rollout_data=None, diffusion_log_images=2), sample)
+
+    assert sample.generated_output is None
+
+
+def test_media_kept_for_logged_samples():
+    sample = make_sample(index=1)
+
+    release_eval_sample_media(Namespace(save_debug_rollout_data=None, diffusion_log_images=2), sample)
+
+    assert sample.generated_output is not None
+
+
+def test_media_kept_when_debug_dump_enabled():
+    sample = make_sample(index=3)
+
+    release_eval_sample_media(Namespace(save_debug_rollout_data="/tmp/dump", diffusion_log_images=0), sample)
+
+    assert sample.generated_output is not None