diff --git a/soup_cli/commands/infer.py b/soup_cli/commands/infer.py index 3336943..217de55 100644 --- a/soup_cli/commands/infer.py +++ b/soup_cli/commands/infer.py @@ -104,24 +104,28 @@ def infer( model_obj, tokenizer = _load_model(str(model_path), base, device) console.print("[green]Model loaded.[/]\n") - # Run inference + # Run inference — stream results to disk as they are generated output_path = Path(output) - results = [] + total_tokens = 0 + num_results = 0 start_time = time.time() - with Progress( - SpinnerColumn(), - TextColumn("[progress.description]{task.description}"), - BarColumn(), - TextColumn("[progress.percentage]{task.percentage:>3.0f}%"), - TimeElapsedColumn(), - console=console, - ) as progress: + with ( + open(output_path, "w", encoding="utf-8") as out_f, + Progress( + SpinnerColumn(), + TextColumn("[progress.description]{task.description}"), + BarColumn(), + TextColumn("[progress.percentage]{task.percentage:>3.0f}%"), + TimeElapsedColumn(), + console=console, + ) as progress, + ): task = progress.add_task("Generating...", total=len(prompts)) - for idx, prompt_text in enumerate(prompts): + for prompt_text in prompts: messages = [{"role": "user", "content": prompt_text}] - response = _generate( + response, token_count = _generate( model_obj, tokenizer, messages, max_tokens=max_tokens, temperature=temperature, ) @@ -129,25 +133,20 @@ def infer( result = { "prompt": prompt_text, "response": response, - "tokens_generated": len(tokenizer.encode(response)) if response else 0, + "tokens_generated": token_count, } - results.append(result) + out_f.write(json.dumps(result, ensure_ascii=False) + "\n") + out_f.flush() + total_tokens += token_count + num_results += 1 progress.update(task, advance=1) elapsed = time.time() - start_time - - # Write results - with open(output_path, "w", encoding="utf-8") as f: - for row in results: - f.write(json.dumps(row, ensure_ascii=False) + "\n") - - # Summary - total_tokens = sum(r["tokens_generated"] for r in results) tokens_per_sec = total_tokens / elapsed if elapsed > 0 else 0 console.print( Panel( - f"Prompts: [bold]{len(results)}[/]\n" + f"Prompts: [bold]{num_results}[/]\n" f"Total tokens: [bold]{total_tokens}[/]\n" f"Duration: [bold]{elapsed:.1f}s[/]\n" f"Throughput: [bold]{tokens_per_sec:.1f} tok/s[/]\n" @@ -227,8 +226,10 @@ def _load_model(model_path: str, base_model: Optional[str], device: str): return model_obj, tokenizer -def _generate(model, tokenizer, messages, max_tokens=256, temperature=0.7) -> str: - """Generate a response from the model.""" +def _generate( + model, tokenizer, messages, max_tokens=256, temperature=0.7, +) -> tuple[str, int]: + """Generate a response from the model. Returns (text, token_count).""" import torch if hasattr(tokenizer, "apply_chat_template") and tokenizer.chat_template: @@ -267,4 +268,6 @@ def _generate(model, tokenizer, messages, max_tokens=256, temperature=0.7) -> st outputs = model.generate(**gen_kwargs) new_tokens = outputs[0][input_ids.shape[1]:] - return tokenizer.decode(new_tokens, skip_special_tokens=True).strip() + token_count = new_tokens.shape[0] + response_text = tokenizer.decode(new_tokens, skip_special_tokens=True).strip() + return response_text, token_count diff --git a/soup_cli/commands/train.py b/soup_cli/commands/train.py index 5853f90..aac3d6b 100644 --- a/soup_cli/commands/train.py +++ b/soup_cli/commands/train.py @@ -92,7 +92,7 @@ def train( # --- TensorBoard setup --- if tensorboard: try: - from torch.utils.tensorboard import SummaryWriter # noqa: F401 + import tensorboard # noqa: F401 console.print("[green]TensorBoard logging enabled[/]") except ImportError: diff --git a/tests/test_tensorboard.py b/tests/test_tensorboard.py index c952384..57dd077 100644 --- a/tests/test_tensorboard.py +++ b/tests/test_tensorboard.py @@ -1,5 +1,6 @@ """Tests for --tensorboard flag in soup train command.""" +from unittest.mock import MagicMock from unittest.mock import patch as mock_patch # ─── Flag Conflict Tests ────────────────────────────────────────────────── @@ -30,6 +31,27 @@ class TestTensorBoardFlagConflict: assert result.exit_code != 0 assert "cannot use" in result.output.lower() or "pick one" in result.output.lower() + def test_tensorboard_only_does_not_conflict(self, tmp_path): + """--tensorboard alone should not trigger conflict error.""" + from typer.testing import CliRunner + + from soup_cli.cli import app + + config_file = tmp_path / "soup.yaml" + config_file.write_text( + "base: some-model\n" + "task: sft\n" + "data:\n" + " train: ./data.jsonl\n" + ) + + runner = CliRunner() + result = runner.invoke(app, [ + "train", "--config", str(config_file), "--tensorboard", + ]) + # Should not hit the conflict error — may fail later (import/data) + assert "cannot use" not in result.output.lower() + # ─── TensorBoard Import Check Tests ────────────────────────────────────── @@ -52,72 +74,22 @@ class TestTensorBoardImportCheck: ) # Mock tensorboard import to fail - with mock_patch.dict("sys.modules", {"torch.utils.tensorboard": None}): + with mock_patch.dict("sys.modules", {"tensorboard": None}): with mock_patch( "builtins.__import__", side_effect=lambda name, *args, **kwargs: ( (_ for _ in ()).throw(ImportError("no tensorboard")) - if "tensorboard" in name else __import__(name, *args, **kwargs) + if name == "tensorboard" else __import__(name, *args, **kwargs) ), ): runner = CliRunner() result = runner.invoke(app, [ "train", "--config", str(config_file), "--tensorboard", ]) - # Either fails with import error or passes to next validation - # The key is it doesn't crash unexpectedly + # Should fail with import error assert result.exit_code != 0 or "tensorboard" in result.output.lower() -# ─── Report To Logic Tests ─────────────────────────────────────────────── - - -class TestReportToLogic: - """Test report_to routing based on flags.""" - - def test_default_report_to_none(self): - """With no flags, report_to should be 'none'.""" - wandb_flag = False - tensorboard_flag = False - - if wandb_flag: - report_to = "wandb" - elif tensorboard_flag: - report_to = "tensorboard" - else: - report_to = "none" - - assert report_to == "none" - - def test_wandb_sets_report_to_wandb(self): - """With --wandb, report_to should be 'wandb'.""" - wandb_flag = True - tensorboard_flag = False - - if wandb_flag: - report_to = "wandb" - elif tensorboard_flag: - report_to = "tensorboard" - else: - report_to = "none" - - assert report_to == "wandb" - - def test_tensorboard_sets_report_to_tensorboard(self): - """With --tensorboard, report_to should be 'tensorboard'.""" - wandb_flag = False - tensorboard_flag = True - - if wandb_flag: - report_to = "wandb" - elif tensorboard_flag: - report_to = "tensorboard" - else: - report_to = "none" - - assert report_to == "tensorboard" - - # ─── CLI Help Tests ────────────────────────────────────────────────────── @@ -146,49 +118,107 @@ class TestTensorBoardCLI: assert "tensorboard" in result.output.lower() -# ─── Integration with Train Routing Tests ──────────────────────────────── +# ─── Trainer report_to Integration Tests ───────────────────────────────── -class TestTensorBoardTrainIntegration: - """Test tensorboard flag integrates with train command routing.""" - - def test_tensorboard_report_to_passed_to_sft_trainer(self, tmp_path): - """When --tensorboard is set, report_to='tensorboard' should be passed.""" +class TestTensorBoardTrainerIntegration: + """Test that trainers receive report_to='tensorboard' when flag is set.""" + def test_sft_trainer_accepts_tensorboard_report_to(self): + """SFT trainer should accept report_to='tensorboard'.""" from soup_cli.config.schema import SoupConfig + from soup_cli.trainer.sft import SFTTrainerWrapper - config = SoupConfig( + cfg = SoupConfig( base="some-model", task="sft", data={"train": "./data.jsonl"}, ) - # Simulate what train command does - report_to = "tensorboard" - assert report_to == "tensorboard" - assert config.task == "sft" + wrapper = SFTTrainerWrapper(cfg, device="cpu", report_to="tensorboard") + assert wrapper.report_to == "tensorboard" - def test_tensorboard_report_to_passed_to_dpo_trainer(self, tmp_path): - """DPO trainer should also accept tensorboard report_to.""" + def test_dpo_trainer_accepts_tensorboard_report_to(self): + """DPO trainer should accept report_to='tensorboard'.""" from soup_cli.config.schema import SoupConfig + from soup_cli.trainer.dpo import DPOTrainerWrapper - config = SoupConfig( + cfg = SoupConfig( base="some-model", task="dpo", data={"train": "./data.jsonl"}, ) - report_to = "tensorboard" - assert report_to == "tensorboard" - assert config.task == "dpo" + wrapper = DPOTrainerWrapper(cfg, device="cpu", report_to="tensorboard") + assert wrapper.report_to == "tensorboard" - def test_tensorboard_report_to_passed_to_grpo_trainer(self, tmp_path): - """GRPO trainer should also accept tensorboard report_to.""" + def test_grpo_trainer_accepts_tensorboard_report_to(self): + """GRPO trainer should accept report_to='tensorboard'.""" from soup_cli.config.schema import SoupConfig + from soup_cli.trainer.grpo import GRPOTrainerWrapper - config = SoupConfig( + cfg = SoupConfig( base="some-model", task="grpo", data={"train": "./data.jsonl"}, ) - report_to = "tensorboard" - assert report_to == "tensorboard" - assert config.task == "grpo" + wrapper = GRPOTrainerWrapper(cfg, device="cpu", report_to="tensorboard") + assert wrapper.report_to == "tensorboard" + + def test_kto_trainer_accepts_tensorboard_report_to(self): + """KTO trainer should accept report_to='tensorboard'.""" + from soup_cli.config.schema import SoupConfig + from soup_cli.trainer.kto import KTOTrainerWrapper + + cfg = SoupConfig( + base="some-model", + task="kto", + data={"train": "./data.jsonl"}, + ) + wrapper = KTOTrainerWrapper(cfg, device="cpu", report_to="tensorboard") + assert wrapper.report_to == "tensorboard" + + +# ─── Sweep Integration Test ────────────────────────────────────────────── + + +class TestTensorBoardSweepRouting: + """Test tensorboard works through sweep routing.""" + + def test_sweep_run_single_with_tensorboard_report_to(self): + """Sweep _run_single should pass report_to through to trainer.""" + from soup_cli.config.schema import SoupConfig + + cfg = SoupConfig( + base="some-model", + task="sft", + data={"train": "./data.jsonl"}, + ) + + fake_dataset = {"train": [{"messages": [{"role": "user", "content": "Q"}]}]} + fake_result = { + "initial_loss": 1.0, + "final_loss": 0.5, + "total_steps": 10, + "duration_secs": 60.0, + "output_dir": "./output", + "duration": "1m", + } + fake_gpu_info = {"memory_total": "0 MB", "memory_total_bytes": 0} + + with mock_patch("soup_cli.data.loader.load_dataset", return_value=fake_dataset), \ + mock_patch("soup_cli.utils.gpu.detect_device", return_value=("cpu", "CPU")), \ + mock_patch("soup_cli.utils.gpu.get_gpu_info", return_value=fake_gpu_info), \ + mock_patch("soup_cli.experiment.tracker.ExperimentTracker") as mock_tracker_cls, \ + mock_patch("soup_cli.monitoring.display.TrainingDisplay"), \ + mock_patch("soup_cli.trainer.sft.SFTTrainerWrapper.setup"), \ + mock_patch( + "soup_cli.trainer.sft.SFTTrainerWrapper.train", return_value=fake_result + ): + mock_tracker = MagicMock() + mock_tracker.start_run.return_value = "run-tb-1" + mock_tracker_cls.return_value = mock_tracker + + from soup_cli.commands.sweep import _run_single + + result = _run_single(cfg, {}, "tb_run_1", None) + + assert result["run_id"] == "run-tb-1"