RomanNum3ral
/
soup
mirror of https://github.com/razor-ai/soup.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix: address code review findings for v0.13.0 

- Fix TensorBoard import check: use `import tensorboard` directly
  (not torch.utils.tensorboard shim) for accurate availability check
- Stream JSONL output during inference instead of buffering in memory
  (crash-safe, handles large prompt files)
- Return accurate token count from _generate via tensor shape instead
  of re-encoding decoded text
- Replace shallow tests with real trainer integration tests that
  verify report_to='tensorboard' is accepted by all trainer wrappers
This commit is contained in:
Alpamys 2026-03-25 18:48:43 +05:00
parent db19f64649
commit 6806cdd944
3 changed files with 136 additions and 103 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

55
soup_cli/commands/infer.py
View File

@ -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

2
soup_cli/commands/train.py
View File

@ -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:

182
tests/test_tensorboard.py
View File

@ -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"