make dataset work

Author: xzrderek

eval_protocol/cli_commands/create_rft.py

@@ -20,6 +20,7 @@
     create_dataset_from_jsonl,
     create_reinforcement_fine_tuning_job,
 )
+from ..fireworks_rft import detect_dataset_builder, materialize_dataset_via_builder
 from .upload import _discover_tests, _normalize_evaluator_id, _prompt_select
 
 
@@ -316,6 +317,9 @@ def create_rft_command(args) -> int:
     non_interactive: bool = bool(getattr(args, "yes", False))
     dry_run: bool = bool(getattr(args, "dry_run", False))
     force: bool = bool(getattr(args, "force", False))
+    # Track the specifically chosen test (if any) to aid dataset inference later
+    selected_test_file_path: Optional[str] = None
+    selected_test_func_name: Optional[str] = None
 
     api_key = get_fireworks_api_key()
     if not api_key:
@@ -354,6 +358,8 @@ def create_rft_command(args) -> int:
         func_name = chosen.qualname.split(".")[-1]
         source_file_name = os.path.splitext(os.path.basename(chosen.file_path))[0]
         evaluator_id = _normalize_evaluator_id(f"{source_file_name}-{func_name}")
+        selected_test_file_path = chosen.file_path
+        selected_test_func_name = func_name
     # Resolve evaluator resource name to fully-qualified format required by API
     evaluator_resource_name = f"accounts/{account_id}/evaluators/{evaluator_id}"
 
@@ -404,6 +410,8 @@ def create_rft_command(args) -> int:
                 except Exception:
                     rel = abs_path
                 selected_entry = f"{rel}::{func_name}"
+                selected_test_file_path = tests[0].file_path
+                selected_test_func_name = func_name
             else:
                 # Try to match evaluator_id to a discovered test's normalized ID
                 for t in tests:
@@ -417,6 +425,8 @@ def create_rft_command(args) -> int:
                         except Exception:
                             rel = abs_path
                         selected_entry = f"{rel}::{func_name}"
+                        selected_test_file_path = t.file_path
+                        selected_test_func_name = func_name
                         break
             # If still unresolved and multiple tests exist, fail fast to avoid uploading unintended evaluators
             if selected_entry is None and len(tests) > 1:
@@ -480,30 +490,48 @@ def create_rft_command(args) -> int:
     dataset_builder = getattr(args, "dataset_builder", None)  # accepted but unused in simplified flow
 
     if not dataset_id:
-        # Prefer explicit --dataset-jsonl, else attempt to extract from data loader or input_dataset of the single discovered test
+        # Prefer explicit --dataset-jsonl, else attempt to extract from the selected test's data loader or input_dataset.
         if not dataset_jsonl:
-            tests = _discover_tests(project_root)
-            if len(tests) == 1:
-                func_name = tests[0].qualname.split(".")[-1]
-                # Try data_loaders first (existing behavior)
-                dataset_jsonl = _extract_jsonl_from_dataloader(tests[0].file_path, func_name)
+            # Use specifically selected test if available; else only infer when exactly one test exists
+            test_file_for_infer = None
+            func_for_infer = None
+            if selected_test_file_path and selected_test_func_name:
+                test_file_for_infer = selected_test_file_path
+                func_for_infer = selected_test_func_name
+            else:
+                tests = _discover_tests(project_root)
+                if len(tests) == 1:
+                    test_file_for_infer = tests[0].file_path
+                    func_for_infer = tests[0].qualname.split(".")[-1]
+            if test_file_for_infer and func_for_infer:
+                # Try data_loaders first
+                dataset_jsonl = _extract_jsonl_from_dataloader(test_file_for_infer, func_for_infer)
                 if dataset_jsonl:
-                    # Display relative path for readability
                     try:
                         rel = os.path.relpath(dataset_jsonl, project_root)
                     except Exception:
                         rel = dataset_jsonl
                     print(f"✓ Using JSONL from data loader: {rel}")
-                else:
+                if not dataset_jsonl:
                     # Fall back to input_dataset (dataset_path)
-                    dataset_jsonl = _extract_jsonl_from_input_dataset(tests[0].file_path, func_name)
+                    dataset_jsonl = _extract_jsonl_from_input_dataset(test_file_for_infer, func_for_infer)
                     if dataset_jsonl:
-                        # Display relative path for readability
                         try:
                             rel = os.path.relpath(dataset_jsonl, project_root)
                         except Exception:
                             rel = dataset_jsonl
                         print(f"✓ Using JSONL from input_dataset: {rel}")
+                if not dataset_jsonl:
+                    # Last resort: attempt to detect and run a dataset builder in the test's directory
+                    metric_dir = os.path.dirname(test_file_for_infer)
+                    builder_spec = detect_dataset_builder(metric_dir)
+                    if builder_spec:
+                        try:
+                            tmp_jsonl, count = materialize_dataset_via_builder(builder_spec)
+                            dataset_jsonl = tmp_jsonl
+                            print(f"✓ Materialized {count} rows via dataset builder: {builder_spec}")
+                        except Exception as e:
+                            print(f"Warning: dataset builder failed: {e}")
         if not dataset_jsonl:
             print(
                 "Error: Could not determine dataset. Provide --dataset-id or --dataset-jsonl, or ensure a JSONL-based data loader or input_dataset is used in your single discovered test."

tests/test_cli_create_rft_infer.py

@@ -386,3 +386,232 @@ def _raise(*a, **k):
 
     rc = cr.create_rft_command(args)
     assert rc == 1
+
+
+def test_create_rft_fallback_to_dataset_builder(tmp_path, monkeypatch):
+    # Setup project
+    project = tmp_path / "proj"
+    project.mkdir()
+    monkeypatch.chdir(project)
+
+    # Single discovered test without data_loaders or input_dataset
+    test_file = project / "metric" / "test_builder.py"
+    test_file.parent.mkdir(parents=True, exist_ok=True)
+    test_file.write_text("# builder case", encoding="utf-8")
+    single_disc = SimpleNamespace(qualname="metric.test_builder", file_path=str(test_file))
+    monkeypatch.setattr(cr, "_discover_tests", lambda cwd: [single_disc])
+
+    # Environment
+    monkeypatch.setenv("FIREWORKS_API_KEY", "fw_dummy")
+    monkeypatch.setenv("FIREWORKS_ACCOUNT_ID", "acct123")
+    monkeypatch.setenv("FIREWORKS_API_BASE", "https://api.fireworks.ai")
+
+    # Stub selector, upload, and polling
+    import eval_protocol.cli_commands.upload as upload_mod
+
+    monkeypatch.setattr(upload_mod, "_prompt_select", lambda tests, non_interactive=False: tests[:1])
+    monkeypatch.setattr(upload_mod, "upload_command", lambda args: 0)
+    monkeypatch.setattr(cr, "_poll_evaluator_status", lambda **kwargs: True)
+
+    # Dataset builder fallback
+    out_jsonl = project / "metric" / "builder_out.jsonl"
+    out_jsonl.write_text('{"row":1}\n{"row":2}\n', encoding="utf-8")
+
+    monkeypatch.setattr(cr, "detect_dataset_builder", lambda metric_dir: "builder.py::build_training_dataset")
+    monkeypatch.setattr(cr, "materialize_dataset_via_builder", lambda spec: (str(out_jsonl), 2))
+
+    # Capture dataset creation args
+    captured = {"dataset_id": None, "jsonl_path": None}
+
+    def _fake_create_dataset_from_jsonl(account_id, api_key, api_base, dataset_id, display_name, jsonl_path):
+        captured["dataset_id"] = dataset_id
+        captured["jsonl_path"] = jsonl_path
+        return dataset_id, {"name": f"accounts/{account_id}/datasets/{dataset_id}", "state": "UPLOADING"}
+
+    monkeypatch.setattr(cr, "create_dataset_from_jsonl", _fake_create_dataset_from_jsonl)
+    monkeypatch.setattr(cr, "create_reinforcement_fine_tuning_job", lambda *a, **k: {"name": "jobs/123"})
+
+    # Run without dataset inputs so builder path is used
+    import argparse
+
+    args = argparse.Namespace(
+        evaluator_id=None,
+        yes=True,
+        dry_run=False,
+        force=False,
+        env_file=None,
+        dataset_id=None,
+        dataset_jsonl=None,
+        dataset_display_name=None,
+        dataset_builder=None,
+        base_model=None,
+        warm_start_from="accounts/acct123/models/ft-abc123",
+        output_model=None,
+        n=None,
+        max_tokens=None,
+        learning_rate=None,
+        batch_size=None,
+        epochs=None,
+        lora_rank=None,
+        max_context_length=None,
+        chunk_size=None,
+        eval_auto_carveout=None,
+    )
+
+    rc = cr.create_rft_command(args)
+    assert rc == 0
+    # Evaluator id derived from test_builder -> "test-builder-test-builder"
+    assert captured["dataset_id"] is not None
+    assert captured["dataset_id"].startswith("test-builder-test-builder-dataset-")
+    # Ensure we used the materialized JSONL
+    assert captured["jsonl_path"] == str(out_jsonl)
+
+
+def test_create_rft_uses_dataloader_jsonl_when_available(tmp_path, monkeypatch):
+    # Setup project
+    project = tmp_path / "proj"
+    project.mkdir()
+    monkeypatch.chdir(project)
+
+    # Single discovered test
+    test_file = project / "metric" / "test_loader.py"
+    test_file.parent.mkdir(parents=True, exist_ok=True)
+    test_file.write_text("# loader case", encoding="utf-8")
+    single_disc = SimpleNamespace(qualname="metric.test_loader", file_path=str(test_file))
+    monkeypatch.setattr(cr, "_discover_tests", lambda cwd: [single_disc])
+
+    # Environment
+    monkeypatch.setenv("FIREWORKS_API_KEY", "fw_dummy")
+    monkeypatch.setenv("FIREWORKS_ACCOUNT_ID", "acct123")
+    monkeypatch.setenv("FIREWORKS_API_BASE", "https://api.fireworks.ai")
+
+    # Stub selector, upload, and polling
+    import eval_protocol.cli_commands.upload as upload_mod
+
+    monkeypatch.setattr(upload_mod, "_prompt_select", lambda tests, non_interactive=False: tests[:1])
+    monkeypatch.setattr(upload_mod, "upload_command", lambda args: 0)
+    monkeypatch.setattr(cr, "_poll_evaluator_status", lambda **kwargs: True)
+
+    # Provide JSONL via dataloader extractor
+    dl_jsonl = project / "metric" / "loader_out.jsonl"
+    dl_jsonl.write_text('{"a":1}\n', encoding="utf-8")
+    monkeypatch.setattr(cr, "_extract_jsonl_from_dataloader", lambda f, fn: str(dl_jsonl))
+    monkeypatch.setattr(cr, "_extract_jsonl_from_input_dataset", lambda f, fn: None)
+    monkeypatch.setattr(cr, "detect_dataset_builder", lambda metric_dir: None)
+
+    captured = {"dataset_id": None, "jsonl_path": None}
+
+    def _fake_create_dataset_from_jsonl(account_id, api_key, api_base, dataset_id, display_name, jsonl_path):
+        captured["dataset_id"] = dataset_id
+        captured["jsonl_path"] = jsonl_path
+        return dataset_id, {"name": f"accounts/{account_id}/datasets/{dataset_id}", "state": "UPLOADING"}
+
+    monkeypatch.setattr(cr, "create_dataset_from_jsonl", _fake_create_dataset_from_jsonl)
+    monkeypatch.setattr(cr, "create_reinforcement_fine_tuning_job", lambda *a, **k: {"name": "jobs/123"})
+
+    import argparse
+
+    args = argparse.Namespace(
+        evaluator_id=None,
+        yes=True,
+        dry_run=False,
+        force=False,
+        env_file=None,
+        dataset_id=None,
+        dataset_jsonl=None,
+        dataset_display_name=None,
+        dataset_builder=None,
+        base_model=None,
+        warm_start_from="accounts/acct123/models/ft-abc123",
+        output_model=None,
+        n=None,
+        max_tokens=None,
+        learning_rate=None,
+        batch_size=None,
+        epochs=None,
+        lora_rank=None,
+        max_context_length=None,
+        chunk_size=None,
+        eval_auto_carveout=None,
+    )
+
+    rc = cr.create_rft_command(args)
+    assert rc == 0
+    assert captured["dataset_id"] is not None
+    assert captured["dataset_id"].startswith("test-loader-test-loader-dataset-")
+    assert captured["jsonl_path"] == str(dl_jsonl)
+
+
+def test_create_rft_uses_input_dataset_jsonl_when_available(tmp_path, monkeypatch):
+    # Setup project
+    project = tmp_path / "proj"
+    project.mkdir()
+    monkeypatch.chdir(project)
+
+    # Single discovered test
+    test_file = project / "metric" / "test_input_ds.py"
+    test_file.parent.mkdir(parents=True, exist_ok=True)
+    test_file.write_text("# input_dataset case", encoding="utf-8")
+    single_disc = SimpleNamespace(qualname="metric.test_input_ds", file_path=str(test_file))
+    monkeypatch.setattr(cr, "_discover_tests", lambda cwd: [single_disc])
+
+    # Environment
+    monkeypatch.setenv("FIREWORKS_API_KEY", "fw_dummy")
+    monkeypatch.setenv("FIREWORKS_ACCOUNT_ID", "acct123")
+    monkeypatch.setenv("FIREWORKS_API_BASE", "https://api.fireworks.ai")
+
+    # Stub selector, upload, and polling
+    import eval_protocol.cli_commands.upload as upload_mod
+
+    monkeypatch.setattr(upload_mod, "_prompt_select", lambda tests, non_interactive=False: tests[:1])
+    monkeypatch.setattr(upload_mod, "upload_command", lambda args: 0)
+    monkeypatch.setattr(cr, "_poll_evaluator_status", lambda **kwargs: True)
+
+    # Provide JSONL via input_dataset extractor
+    id_jsonl = project / "metric" / "input_ds_out.jsonl"
+    id_jsonl.write_text('{"b":2}\n', encoding="utf-8")
+    monkeypatch.setattr(cr, "_extract_jsonl_from_dataloader", lambda f, fn: None)
+    monkeypatch.setattr(cr, "_extract_jsonl_from_input_dataset", lambda f, fn: str(id_jsonl))
+    monkeypatch.setattr(cr, "detect_dataset_builder", lambda metric_dir: None)
+
+    captured = {"dataset_id": None, "jsonl_path": None}
+
+    def _fake_create_dataset_from_jsonl(account_id, api_key, api_base, dataset_id, display_name, jsonl_path):
+        captured["dataset_id"] = dataset_id
+        captured["jsonl_path"] = jsonl_path
+        return dataset_id, {"name": f"accounts/{account_id}/datasets/{dataset_id}", "state": "UPLOADING"}
+
+    monkeypatch.setattr(cr, "create_dataset_from_jsonl", _fake_create_dataset_from_jsonl)
+    monkeypatch.setattr(cr, "create_reinforcement_fine_tuning_job", lambda *a, **k: {"name": "jobs/123"})
+
+    import argparse
+
+    args = argparse.Namespace(
+        evaluator_id=None,
+        yes=True,
+        dry_run=False,
+        force=False,
+        env_file=None,
+        dataset_id=None,
+        dataset_jsonl=None,
+        dataset_display_name=None,
+        dataset_builder=None,
+        base_model=None,
+        warm_start_from="accounts/acct123/models/ft-abc123",
+        output_model=None,
+        n=None,
+        max_tokens=None,
+        learning_rate=None,
+        batch_size=None,
+        epochs=None,
+        lora_rank=None,
+        max_context_length=None,
+        chunk_size=None,
+        eval_auto_carveout=None,
+    )
+
+    rc = cr.create_rft_command(args)
+    assert rc == 0
+    assert captured["dataset_id"] is not None
+    assert captured["dataset_id"].startswith("test-input-ds-test-input-ds-dataset-")
+    assert captured["jsonl_path"] == str(id_jsonl)