Stop saving prepared datasets in outputs

docs/reference/workspace-structure.md

@@ -42,15 +42,15 @@ datasets/my_dataset/
 ```
 
 Each unprepared `labels.csv` must include `id` and `label` columns. Your source
-`datasets/<name>/` folder is mounted read-only and is never modified. When a run
-starts, Agentomics copies its public splits into the run workspace and converts
-their labels to `id,numeric_label`. The `input/` interface is recorded at
-preparation time, must match across all splits, and must not be modified during
-a run. The optional source `test/` split is excluded from the agent worker's
-mounts and remains outside the agent-facing prepared data.
+`datasets/<name>/` folder is mounted read-only and is never modified. While a
+run is active, Agentomics creates a transient prepared view in the run workspace
+and converts its labels to `id,numeric_label`. The `input/` interface is
+recorded at preparation time, must match across all splits, and must not be
+modified during a run. The optional source `test/` split is excluded from the
+agent worker's mounts and remains outside the agent-facing prepared data.
 
-The prepared, agent-facing splits are written to the run workspace (never back
-into `datasets/`):
+Exact versioned splits produced or selected by the run are persisted in the run
+workspace (never back into the original `datasets/` folder):
 
 ```text
 outputs/<agent_id>/run/shared/splits/split_0/
@@ -73,6 +73,16 @@ datasets/my_dataset/test/
 └── labels.csv          # id,label
 ```
 
+The transient `prepared_datasets/` directory, including any intermediate CSV
+conversion data inside it, is removed after test evaluation and report
+generation. It is regenerated from the original dataset when a run is forked
+and is not part of the final output.
+
+The resolved preparation choices remain in
+`run/shared/dataset_metadata.json`. This small file is persistent run state, so
+a fork can recreate the prepared dataset without asking again for task type,
+label mapping, or the CSV label column.
+
 ## Active Workspace
 
 The active host workspace is mounted at `/workspace` in the container.
@@ -85,6 +95,7 @@ Current run working directory:
 <workspace_root>/run/
 ├── shared/
 │   ├── config.json
+│   ├── dataset_metadata.json
 │   ├── environment.yml
 │   └── splits/
 ├── current_iteration/
@@ -205,7 +216,7 @@ outputs/<agent_id>/
 ### run/
 
 The working directory. `shared/` holds the run config, the shared conda
-environment, the prepared dataset copy, and the train/validation splits.
+environment, and the persisted train/validation splits.
 `current_iteration/` is the active iteration while the run is in progress;
 completed iterations are archived as `iteration_N/`.
 

docs/user-guide/forking.md

@@ -73,17 +73,27 @@ Dataset and validation metric are locked to keep all iterations comparable acros
 
 When a fork is set up, the following happens before the new run starts:
 
-1. The source workspace state is copied, excluding generated reports/logs and untracked Conda environments.
+1. The source workspace state is copied, excluding generated reports/logs,
+   untracked Conda environments, and transient dataset-preparation directories.
 2. The git history in the run directory is checked out at the requested checkpoint — files added in later commits are removed.
 3. Absolute paths stored in step outputs are rewritten to point to the new workspace.
 4. The shared Conda environment is rebuilt from the checkpoint's `environment.yml` using the new run ID.
+5. The source dataset is prepared again at the same container path before the
+   fork resumes. Exact versioned splits from the checkpoint are retained rather
+   than regenerated.
 
 The forked run then continues from that state exactly as if the original run had stopped there.
 
-**Note on supplementary materials**: Forked runs reference the same dataset directory as the source run. Any changes to dataset files (including `supplementary/`) will be visible to the fork.
+The preparation uses `run/shared/dataset_metadata.json` from the selected
+checkpoint. Resolved choices such as task type, label mapping, and the CSV label
+column therefore remain unchanged and are not prompted for again.
+
+The original dataset must remain available under `--datasets-dir` with the same
+dataset name. Forked runs reference that dataset directory, so supplementary
+materials are available without being copied into every run output.
 
 Split directories may contain symbolic links pointing to container paths. Forks
-run in the same Docker layout, so those links continue to resolve.
+prepare the dataset at the same Docker path, so those links continue to resolve.
 
 ## Example: Extend a Completed Run
 

docs/user-guide/outputs.md

@@ -23,6 +23,7 @@ outputs/<agent_id>/
 ├── run/                      # All iterations + shared run state
 │   ├── shared/
 │   │   ├── config.json
+│   │   ├── dataset_metadata.json
 │   │   └── splits/
 │   ├── iteration_0/
 │   ├── iteration_1/
@@ -70,6 +71,10 @@ agentomics-inference --agent-dir outputs/<agent_id> --input data/input --output
 
 ## Iteration Directories
 
+`run/shared/dataset_metadata.json` stores the resolved dataset preparation
+metadata. Forks reuse it so task type, label mapping, and the CSV label column
+are not requested again.
+
 Each iteration's files are preserved under `run/iteration_N/`:
 
 ```
@@ -120,6 +125,9 @@ separate staging area or temporary volume:
 
 - The host workspace defaults to `outputs/<agent_id>/` and is mounted at
   `/workspace` in the container.
+- Dataset preparation directories may exist there while the run is active, but
+  they are removed after test evaluation and reporting. Exact versioned splits
+  remain under `run/shared/splits/`.
 
 ## W&B Logging
 

src/agentomics/cli/run.py

@@ -14,14 +14,15 @@
 )
 from agentomics.cli.inference import run_inference_in_docker
 from agentomics.cli.run_arguments import add_run_arguments
-from agentomics.datasets.data_contract import TEST_SPLIT_PREFIX
+from agentomics.datasets.data_contract import PREPARED_DATASETS_DIR_NAME, TEST_SPLIT_PREFIX
 from agentomics.datasets.dataset_preparation import prepare_test_dataset
 from agentomics.datasets.datasets_interactive import (
     get_all_datasets_info,
     interactive_dataset_selection,
     print_datasets_table,
 )
 from agentomics.runtime.read_write_utils import load_config_from_run_dir
+from agentomics.runtime.filesystem import remove_path
 from agentomics.utils.agent_id import create_agent_id
 from agentomics.utils.config import Config
 
@@ -297,7 +298,7 @@ def _run_test_evaluation_in_docker(
 
     metadata_path = (
         workspace_directory
-        / "prepared_datasets"
+        / PREPARED_DATASETS_DIR_NAME
         / dataset_directory.name
         / "metadata.json"
     )
@@ -388,27 +389,31 @@ def main() -> int:
     agent_id = create_agent_id()
     workspace_directory = _resolve_workspace(arguments.workspace_dir, agent_id)
     dataset_mounts = _build_dataset_mounts(dataset_directory)
-    exit_code = _run_agent_in_docker(
-        arguments,
-        workspace_directory,
-        agent_id,
-        dataset_mounts,
-    )
-    if _is_agent_run(arguments) and exit_code == 0:
-        try:
-            _run_test_evaluation_in_docker(
-                image=arguments.image,
-                cpu_only=arguments.cpu_only,
-                workspace_directory=workspace_directory,
-                dataset_directory=dataset_directory,
-            )
-        finally:
-            _run_reporting_in_docker(
-                arguments.image,
-                workspace_directory,
-                agent_id,
-            )
-    return exit_code
+    try:
+        exit_code = _run_agent_in_docker(
+            arguments,
+            workspace_directory,
+            agent_id,
+            dataset_mounts,
+        )
+        if _is_agent_run(arguments) and exit_code == 0:
+            try:
+                _run_test_evaluation_in_docker(
+                    image=arguments.image,
+                    cpu_only=arguments.cpu_only,
+                    workspace_directory=workspace_directory,
+                    dataset_directory=dataset_directory,
+                )
+            finally:
+                _run_reporting_in_docker(
+                    arguments.image,
+                    workspace_directory,
+                    agent_id,
+                )
+        return exit_code
+    finally:
+        if _is_agent_run(arguments):
+            remove_path(workspace_directory / PREPARED_DATASETS_DIR_NAME)
 
 
 if __name__ == "__main__":

src/agentomics/datasets/data_contract.py

@@ -11,6 +11,8 @@
 
 INPUT_DIR_NAME = "input"
 SUPPLEMENTARY_DIR_NAME = "supplementary"
+PREPARED_DATASETS_DIR_NAME = "prepared_datasets"
+DATASET_METADATA_FILE_NAME = "dataset_metadata.json"
 LABELS_FILE_NAME = "labels.csv"
 ID_COLUMN_NAME = "id"
 LABEL_COLUMN_NAME = "label"

src/agentomics/run_agent.py

@@ -7,7 +7,11 @@
 
 from agentomics.run_logging.wandb_setup import setup_logging
 from agentomics.runtime.git_checkpoints import initialize_repo_if_needed
-from agentomics.runtime.read_write_utils import initialize_run_directories, load_dataset_metadata, save_config
+from agentomics.runtime.read_write_utils import (
+    initialize_run_directories,
+    save_config,
+    save_dataset_metadata,
+)
 from agentomics.runtime.run_lifecycle import run_agentomics
 from agentomics.utils.config import Config
 from agentomics.run_logging.env_utils import are_wandb_vars_available
@@ -33,6 +37,7 @@ async def run_experiment(
     split_allowed_iterations: int,
     exploration_iterations: int,
     input_structure: list[str],
+    dataset_metadata: dict,
     label_to_scalar: dict[str, int] | None = None,
     disable_training_reporting: bool = False,
     conda_export_mode: str = "full"
@@ -73,6 +78,7 @@ async def run_experiment(
     initialize_run_directories(config)
     config.wandb_run_id = setup_logging(config) if are_wandb_vars_available() else None
     save_config(config)
+    save_dataset_metadata(config, dataset_metadata)
     initialize_repo_if_needed(config)
 
     config.print_summary()

src/agentomics/run_agent_interactive.py

@@ -2,12 +2,14 @@
 
 import argparse
 import asyncio
-import json
 import sys
 
 from rich.console import Console
 
-from agentomics.datasets.data_contract import METADATA_FILE_NAME, VALIDATION_SPLIT
+from agentomics.datasets.data_contract import (
+    PREPARED_DATASETS_DIR_NAME,
+    VALIDATION_SPLIT,
+)
 from agentomics.datasets.dataset_preparation import prepare_dataset
 from agentomics.datasets.datasets_interactive import (
     get_all_datasets_info,
@@ -16,7 +18,11 @@
 )
 from agentomics.run_agent import run_experiment
 from agentomics.run_logging.env_utils import are_wandb_vars_available
-from agentomics.runtime.read_write_utils import get_next_iteration_index, load_config_from_run_dir
+from agentomics.runtime.read_write_utils import (
+    get_next_iteration_index,
+    load_config_from_run_dir,
+    load_dataset_metadata,
+)
 from agentomics.utils.config import Config
 from agentomics.utils.metrics import resolve_val_metric
 from agentomics.utils.metrics_interactive import display_metrics_table
@@ -28,7 +34,7 @@
 console = Console()
 
 
-def resolve_run_arguments(arguments: argparse.Namespace) -> None:
+def resolve_run_arguments(arguments: argparse.Namespace) -> Config | None:
     existing_config = load_config_from_run_dir(
         arguments.workspace_dir / Config.RUN_DIRNAME,
         missing_ok=True,
@@ -93,6 +99,7 @@ def resolve_run_arguments(arguments: argparse.Namespace) -> None:
         arguments.run_python_timeout = Config.DEFAULT_RUN_PYTHON_TOOL_TIMEOUT
     if arguments.user_prompt is None:
         arguments.user_prompt = Config.DEFAULT_USER_PROMPT
+    return existing_config
 
 
 def _is_tty_available() -> bool:
@@ -137,7 +144,7 @@ def _resolve_interactive_parameters(
 
 
 def run_agent_interactive(arguments: argparse.Namespace) -> int:
-    resolve_run_arguments(arguments)
+    existing_config = resolve_run_arguments(arguments)
 
     if arguments.list_datasets:
         console.print("Available Datasets", style="cyan")
@@ -165,18 +172,23 @@ def run_agent_interactive(arguments: argparse.Namespace) -> int:
 
     dataset, model, iterations = _resolve_interactive_parameters(arguments, provider)
     iteration_plan_model = arguments.iteration_plan_model or model
-    prepared_datasets_dir = arguments.workspace_dir / "prepared_datasets"
-    prepared_dataset_dir = prepared_datasets_dir / dataset
-    if arguments.fork_from_run is not None:
-        metadata_path = prepared_dataset_dir / METADATA_FILE_NAME
-        dataset_metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
-    else:
-        dataset_metadata = prepare_dataset(
-            source_dir=arguments.datasets_dir / dataset,
-            destination_dir=prepared_dataset_dir,
-            task_type=arguments.task_type,
-            interactive=_is_tty_available(),
-        )
+    prepared_datasets_dir = arguments.workspace_dir / PREPARED_DATASETS_DIR_NAME
+    inherited_dataset_metadata = (
+        load_dataset_metadata(existing_config)
+        if existing_config is not None
+        else {}
+    )
+    dataset_metadata = prepare_dataset(
+        source_dir=arguments.datasets_dir / dataset,
+        destination_dir=prepared_datasets_dir / dataset,
+        task_type=(
+            inherited_dataset_metadata.get("task_type")
+            or arguments.task_type
+        ),
+        label_to_scalar=inherited_dataset_metadata.get("label_to_scalar"),
+        label_column=inherited_dataset_metadata.get("label_column"),
+        interactive=_is_tty_available() and existing_config is None,
+    )
     task_type = dataset_metadata["task_type"]
     val_metric = resolve_val_metric(task_type, arguments.val_metric)
     split_allowed_iterations = arguments.split_allowed_iterations
@@ -189,6 +201,7 @@ def run_agent_interactive(arguments: argparse.Namespace) -> int:
             iteration_plan_model=iteration_plan_model,
             dataset_name=dataset,
             task_type=task_type,
+            dataset_metadata=dataset_metadata,
             label_to_scalar=dataset_metadata.get("label_to_scalar"),
             input_structure=dataset_metadata["input_structure"],
             val_metric=val_metric,

src/agentomics/runtime/git_checkpoints.py

@@ -3,6 +3,7 @@
 import subprocess
 from pathlib import Path
 
+from agentomics.datasets.data_contract import PREPARED_DATASETS_DIR_NAME
 from agentomics.runtime.conda_utils import ENVIRONMENT_ARCHIVE_FILENAME
 from agentomics.utils.config import Config
 
@@ -99,6 +100,7 @@ def _write_gitignore(config: Config) -> None:
         "run/shared/.conda/",
         f"{Config.BEST_ITERATION_SNAPSHOT_DIRNAME}/.conda/",
         f"{Config.BEST_ITERATION_SNAPSHOT_DIRNAME}/runtime_info/{ENVIRONMENT_ARCHIVE_FILENAME}",
+        f"{PREPARED_DATASETS_DIR_NAME}/",
         "__pycache__/",
         ".cache/",
         "*.pyc",

src/agentomics/runtime/read_write_utils.py

@@ -8,6 +8,7 @@
 from pathlib import Path
 from typing import Any
 
+from agentomics.datasets.data_contract import DATASET_METADATA_FILE_NAME
 from agentomics.runtime.conda_utils import export_shared_environment_descriptor
 from agentomics.runtime.filesystem import (
     chown_tree_to_root,
@@ -29,6 +30,14 @@ def save_config(config: Config) -> None:
     config.config_path.parent.mkdir(parents=True, exist_ok=True)
     config.config_path.write_text(json.dumps(asdict(config), indent=2), encoding="utf-8")
 
+def save_dataset_metadata(config: Config, metadata: dict[str, Any]) -> None:
+    metadata_path = config.shared_dir / DATASET_METADATA_FILE_NAME
+    metadata_path.parent.mkdir(parents=True, exist_ok=True)
+    metadata_path.write_text(
+        json.dumps(metadata, indent=2),
+        encoding="utf-8",
+    )
+
 def load_config(config_path: Path | str, missing_ok: bool = False) -> Config | None:
     config_path = Path(config_path)
     if not config_path.exists():
@@ -209,8 +218,9 @@ def update_current_iteration_state(config: Config, **changes: object) -> None:
     iteration_state.update(changes)
     _write_json(iteration_state_path, iteration_state)
 
-def load_dataset_metadata(config: Config) -> dict[str, str]:
-    return json.loads((config.dataset_dir / "metadata.json").read_text(encoding="utf-8"))
+def load_dataset_metadata(config: Config) -> dict[str, Any]:
+    dataset_metadata_path = config.shared_dir / DATASET_METADATA_FILE_NAME
+    return json.loads(dataset_metadata_path.read_text(encoding="utf-8"))
 
 def replace_string_in_tree_files(root_dir: Path, old: str, new: str, skip_dirs: set[str] | None = None) -> None:
     for file_path in root_dir.rglob("*"):