From 3c5c6f5846157bfd1d903a3c7fa2ffb0292ad1d1 Mon Sep 17 00:00:00 2001 From: Andrea Gariboldi Date: Thu, 30 Jul 2026 17:10:25 +0200 Subject: [PATCH 1/3] Save dataset metadata and not prepared dataset in output folder --- src/agentomics/cli/run.py | 51 ++++++++++++---------- src/agentomics/datasets/data_contract.py | 2 + src/agentomics/run_agent.py | 8 +++- src/agentomics/run_agent_interactive.py | 47 ++++++++++++-------- src/agentomics/runtime/git_checkpoints.py | 2 + src/agentomics/runtime/read_write_utils.py | 14 +++++- src/agentomics/runtime/setup_fork.py | 8 +++- 7 files changed, 88 insertions(+), 44 deletions(-) diff --git a/src/agentomics/cli/run.py b/src/agentomics/cli/run.py index 30cca7ee..dabf2342 100644 --- a/src/agentomics/cli/run.py +++ b/src/agentomics/cli/run.py @@ -14,7 +14,7 @@ ) 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, @@ -22,6 +22,7 @@ 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__": diff --git a/src/agentomics/datasets/data_contract.py b/src/agentomics/datasets/data_contract.py index c50c90e2..9050a278 100644 --- a/src/agentomics/datasets/data_contract.py +++ b/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" diff --git a/src/agentomics/run_agent.py b/src/agentomics/run_agent.py index c9cad970..0093987d 100644 --- a/src/agentomics/run_agent.py +++ b/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() diff --git a/src/agentomics/run_agent_interactive.py b/src/agentomics/run_agent_interactive.py index 85179b4e..137d4a96 100644 --- a/src/agentomics/run_agent_interactive.py +++ b/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, diff --git a/src/agentomics/runtime/git_checkpoints.py b/src/agentomics/runtime/git_checkpoints.py index a8427ea1..14c18027 100644 --- a/src/agentomics/runtime/git_checkpoints.py +++ b/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", diff --git a/src/agentomics/runtime/read_write_utils.py b/src/agentomics/runtime/read_write_utils.py index 94b24ef1..3eb7a7f4 100644 --- a/src/agentomics/runtime/read_write_utils.py +++ b/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("*"): diff --git a/src/agentomics/runtime/setup_fork.py b/src/agentomics/runtime/setup_fork.py index ff75f39e..98b7172e 100644 --- a/src/agentomics/runtime/setup_fork.py +++ b/src/agentomics/runtime/setup_fork.py @@ -6,6 +6,7 @@ from dataclasses import replace from pathlib import Path +from agentomics.datasets.data_contract import PREPARED_DATASETS_DIR_NAME from agentomics.runtime.conda_utils import ( ensure_environment_from_descriptor, get_shared_environment_path, @@ -40,7 +41,12 @@ def fork_run( target_workspace_dir, symlinks=True, copy_function=shutil.copy2, - ignore=shutil.ignore_patterns(".conda", "reports", "logs"), + ignore=shutil.ignore_patterns( + ".conda", + "reports", + "logs", + PREPARED_DATASETS_DIR_NAME, + ), dirs_exist_ok=True, ) From a37da1eb5a04bc8302fbac290fddf565587d8d49 Mon Sep 17 00:00:00 2001 From: Andrea Gariboldi Date: Thu, 30 Jul 2026 17:10:52 +0200 Subject: [PATCH 2/3] Update docs --- docs/reference/workspace-structure.md | 29 ++++++++++++++++++--------- docs/user-guide/forking.md | 16 ++++++++++++--- docs/user-guide/outputs.md | 8 ++++++++ 3 files changed, 41 insertions(+), 12 deletions(-) diff --git a/docs/reference/workspace-structure.md b/docs/reference/workspace-structure.md index e364352d..5a488ff4 100644 --- a/docs/reference/workspace-structure.md +++ b/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//` 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//` 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//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: /run/ ├── shared/ │ ├── config.json +│ ├── dataset_metadata.json │ ├── environment.yml │ └── splits/ ├── current_iteration/ @@ -205,7 +216,7 @@ outputs// ### 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/`. diff --git a/docs/user-guide/forking.md b/docs/user-guide/forking.md index f5eef9f8..4a485471 100644 --- a/docs/user-guide/forking.md +++ b/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 diff --git a/docs/user-guide/outputs.md b/docs/user-guide/outputs.md index 9ad3d8da..b0e71008 100644 --- a/docs/user-guide/outputs.md +++ b/docs/user-guide/outputs.md @@ -23,6 +23,7 @@ outputs// ├── 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/ --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//` 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 From c1e4b3c68e94a3d93fd4e3398cd23e2381e8a0e4 Mon Sep 17 00:00:00 2001 From: Andrea Gariboldi Date: Thu, 30 Jul 2026 17:11:04 +0200 Subject: [PATCH 3/3] Add tests --- test/test_csv_converter.py | 48 +++++++++ test/test_forking.py | 92 +++++++++++++++- test/test_transient_dataset_state.py | 153 +++++++++++++++++++++++++++ 3 files changed, 292 insertions(+), 1 deletion(-) create mode 100644 test/test_transient_dataset_state.py diff --git a/test/test_csv_converter.py b/test/test_csv_converter.py index 665cd0d2..d9b5b826 100644 --- a/test/test_csv_converter.py +++ b/test/test_csv_converter.py @@ -1,3 +1,4 @@ +import json import sys import unittest from pathlib import Path @@ -11,6 +12,7 @@ sys.path.insert(0, str(SRC_PATH)) from agentomics.datasets.csv_converter import convert_csv_dataset +from agentomics.datasets.dataset_preparation import prepare_dataset class CsvConverterTest(unittest.TestCase): @@ -61,6 +63,52 @@ def test_train_split_is_required(self): self.assertIn("train", str(raised.exception)) + def test_preparation_can_replay_resolved_csv_metadata(self): + with TemporaryDirectory() as tmp: + root = Path(tmp) + source_dir = root / "datasets" / "toy" + source_dir.mkdir(parents=True) + (source_dir / "metadata.json").write_text( + json.dumps({"id_column": "sample_id"}), + encoding="utf-8", + ) + pd.DataFrame( + { + "sample_id": ["row-1", "row-2"], + "feature": [1, 2], + "outcome": ["no", "yes"], + } + ).to_csv(source_dir / "train.csv", index=False) + + first_metadata = prepare_dataset( + source_dir=source_dir, + destination_dir=root / "first_prepared", + task_type="classification", + label_to_scalar={"no": 0, "yes": 1}, + label_column="outcome", + interactive=False, + ) + second_metadata = prepare_dataset( + source_dir=source_dir, + destination_dir=root / "second_prepared", + task_type=first_metadata["task_type"], + label_to_scalar=first_metadata["label_to_scalar"], + label_column=first_metadata["label_column"], + interactive=False, + ) + + self.assertEqual(second_metadata["label_column"], "outcome") + self.assertEqual(second_metadata["id_column"], "sample_id") + self.assertEqual( + second_metadata["label_to_scalar"], + {"no": 0, "yes": 1}, + ) + replayed_labels = pd.read_csv( + root / "second_prepared" / "train" / "labels.csv", + dtype={"id": str}, + ) + self.assertEqual(replayed_labels["id"].tolist(), ["row-1", "row-2"]) + if __name__ == "__main__": unittest.main() diff --git a/test/test_forking.py b/test/test_forking.py index c6a9f4ab..f34655e3 100644 --- a/test/test_forking.py +++ b/test/test_forking.py @@ -15,13 +15,20 @@ sys.path.insert(0, str(SRC_PATH)) from agentomics.runtime.setup_fork import fork_run +from agentomics.datasets.data_contract import PREPARED_DATASETS_DIR_NAME from agentomics.runtime.read_write_utils import replace_string_in_tree_files from agentomics.runtime.git_checkpoints import ( build_iteration_end_commit_message, build_step_commit_message, _find_checkpoint_commit, + initialize_repo_if_needed, +) +from agentomics.runtime.read_write_utils import ( + initialize_run_directories, + load_dataset_metadata, + save_config, + save_dataset_metadata, ) -from agentomics.runtime.read_write_utils import initialize_run_directories, save_config from agentomics.utils.config import Config @@ -167,6 +174,14 @@ def _build_source_run(self, run_id: str) -> None: ) initialize_run_directories(source_config) save_config(source_config) + save_dataset_metadata( + source_config, + { + "task_type": "classification", + "label_to_scalar": {"negative": 0, "positive": 1}, + "input_structure": ["data.csv"], + }, + ) _setup_git_repo(self.source_workspace) run_dir = source_config.run_dir @@ -276,6 +291,60 @@ def test_rebuilds_shared_environment_on_container_local_storage(self): Path(f"/tmp/agentomics/envs/{target_run_id}_env"), ) + def test_does_not_copy_transient_dataset_state(self): + self._build_source_run("src_run") + marker = ( + self.source_workspace + / PREPARED_DATASETS_DIR_NAME + / "_csv_converted_source" + / "toy" + / "marker.txt" + ) + marker.parent.mkdir(parents=True, exist_ok=True) + marker.write_text("transient", encoding="utf-8") + + with patch("agentomics.runtime.setup_fork.ensure_environment_from_descriptor"): + fork_run( + source_workspace_dir=self.source_workspace, + target_agent_id="tgt_run", + target_workspace_dir=self.target_workspace, + fork_from_step=None, + fork_from_iteration=None, + ) + + self.assertFalse( + (self.target_workspace / PREPARED_DATASETS_DIR_NAME).exists() + ) + + def test_preserves_dataset_metadata_from_selected_checkpoint(self): + self._fork( + "src_run", + "tgt_run", + fork_from_step="data_split", + fork_from_iteration=0, + ) + target_config = Config( + agent_id="tgt_run", + model_name="test-model", + iteration_plan_model_name="test-model", + dataset="toy", + tags=[], + val_metric="ACC", + workspace_dir=str(self.target_workspace), + datasets_dir=str(self.root / "datasets"), + task_type="classification", + input_structure=["data.csv"], + ) + + self.assertEqual( + load_dataset_metadata(target_config), + { + "task_type": "classification", + "label_to_scalar": {"negative": 0, "positive": 1}, + "input_structure": ["data.csv"], + }, + ) + def test_raises_when_target_workspace_not_empty(self): self._build_source_run("src_run") self.target_workspace.mkdir(parents=True) @@ -398,6 +467,27 @@ def test_handles_tracked_file_modified_after_last_commit(self): fork_from_iteration=0, ) + def test_transient_dataset_directories_are_ignored(self): + config = Config( + agent_id="gitignore_run", + model_name="test-model", + iteration_plan_model_name="test-model", + dataset="toy", + tags=[], + val_metric="ACC", + workspace_dir=str(self.source_workspace), + datasets_dir=str(self.root / "datasets"), + task_type="classification", + input_structure=["data.csv"], + ) + initialize_run_directories(config) + initialize_repo_if_needed(config) + + gitignore = (self.source_workspace / ".gitignore").read_text( + encoding="utf-8" + ) + self.assertIn(f"{PREPARED_DATASETS_DIR_NAME}/", gitignore) + if __name__ == "__main__": unittest.main() diff --git a/test/test_transient_dataset_state.py b/test/test_transient_dataset_state.py new file mode 100644 index 00000000..acc44872 --- /dev/null +++ b/test/test_transient_dataset_state.py @@ -0,0 +1,153 @@ +from __future__ import annotations + +import argparse +import sys +import tempfile +import unittest +from pathlib import Path +from unittest.mock import Mock, patch + +REPO_ROOT = Path(__file__).resolve().parents[1] +SRC_PATH = REPO_ROOT / "src" +if str(SRC_PATH) not in sys.path: + sys.path.insert(0, str(SRC_PATH)) + +from agentomics.cli import run as run_cli +from agentomics.datasets.data_contract import ( + DATASET_METADATA_FILE_NAME, + PREPARED_DATASETS_DIR_NAME, +) +from agentomics.runtime.filesystem import remove_path +from agentomics.runtime.read_write_utils import ( + load_dataset_metadata, + save_dataset_metadata, +) +from agentomics.utils.config import Config + + +class TestRemoveTransientDatasetState(unittest.TestCase): + def test_removes_transient_directories_but_preserves_splits(self): + with tempfile.TemporaryDirectory() as temporary_directory: + workspace = Path(temporary_directory) + persistent_split = workspace / "run" / "shared" / "splits" / "split_0" + persistent_split.mkdir(parents=True) + (persistent_split / "labels.csv").write_text( + "id,numeric_label\nrow-1,0\n", + encoding="utf-8", + ) + transient_file = ( + workspace + / PREPARED_DATASETS_DIR_NAME + / "_csv_converted_source" + / "toy" + / "data.csv" + ) + transient_file.parent.mkdir(parents=True) + transient_file.write_text("data", encoding="utf-8") + + remove_path(workspace / PREPARED_DATASETS_DIR_NAME) + + self.assertFalse( + (workspace / PREPARED_DATASETS_DIR_NAME).exists() + ) + self.assertTrue((persistent_split / "labels.csv").is_file()) + + +class TestDatasetMetadataPersistence(unittest.TestCase): + def test_metadata_is_stored_in_run_shared_independently_of_prepared_dataset(self): + with tempfile.TemporaryDirectory() as temporary_directory: + workspace = Path(temporary_directory) + config = Mock( + shared_dir=( + workspace / Config.RUN_DIRNAME / Config.SHARED_DIRNAME + ) + ) + metadata = { + "task_type": "classification", + "label_to_scalar": {"negative": 0, "positive": 1}, + "label_column": "outcome", + "id_column": "sample_id", + "input_structure": ["data.csv"], + } + + save_dataset_metadata(config, metadata) + self.assertTrue( + (config.shared_dir / DATASET_METADATA_FILE_NAME).is_file() + ) + prepared_datasets_dir = workspace / PREPARED_DATASETS_DIR_NAME + prepared_datasets_dir.mkdir() + remove_path(prepared_datasets_dir) + + self.assertEqual(load_dataset_metadata(config), metadata) + +class TestTransientDatasetCleanupLifecycle(unittest.TestCase): + def test_cleanup_runs_after_test_evaluation_and_reporting(self): + with tempfile.TemporaryDirectory() as temporary_directory: + root = Path(temporary_directory) + datasets_dir = root / "datasets" + dataset_dir = datasets_dir / "toy" + dataset_dir.mkdir(parents=True) + workspace = root / "workspace" + arguments = argparse.Namespace( + datasets_dir=datasets_dir, + list_datasets=False, + cpu_only=True, + image="test-image", + workspace_dir=workspace, + fork_from_run=None, + dataset="toy", + test=False, + list_models=False, + list_metrics=False, + ) + parser = Mock() + parser.parse_args.return_value = arguments + events = [] + + def run_agent(*_args): + path = ( + workspace + / PREPARED_DATASETS_DIR_NAME + / "_csv_converted_source" + / "toy" + ) + path.mkdir(parents=True) + events.append("run") + return 0 + + def assert_transient_state_exists(event): + self.assertTrue( + (workspace / PREPARED_DATASETS_DIR_NAME).is_dir() + ) + events.append(event) + + with ( + patch.object(run_cli, "build_parser", return_value=parser), + patch.object(run_cli, "create_agent_id", return_value="new_run"), + patch.object(run_cli, "_run_agent_in_docker", side_effect=run_agent), + patch.object( + run_cli, + "_run_test_evaluation_in_docker", + side_effect=lambda **_kwargs: assert_transient_state_exists( + "test evaluation" + ), + ), + patch.object( + run_cli, + "_run_reporting_in_docker", + side_effect=lambda *_args: assert_transient_state_exists( + "reporting" + ), + ), + ): + exit_code = run_cli.main() + + self.assertEqual(exit_code, 0) + self.assertEqual(events, ["run", "test evaluation", "reporting"]) + self.assertFalse( + (workspace / PREPARED_DATASETS_DIR_NAME).exists() + ) + + +if __name__ == "__main__": + unittest.main()