feat(cli): add batch size option

[sarvesh1327]

Summary

• Add a --batch-size CLI option for dynamic KV cache footprint calculations
• Pass the parsed batch size through analyze_model into calculate_footprint
• Cover the new flag default/override and calculator forwarding path

Fixes #2.

Test Plan

• PYTHONPATH=src python3 -m pytest tests/test_cli.py tests/test_calculator.py -q
• PYTHONPATH=src python3 -m pytest -q
• PYTHONPATH=src python3 -m compileall -q src tests
• PYTHONPATH=src python3 -m ruff check src/modelinfo/cli.py tests/test_cli.py

Notes

• Scope intentionally small.

Summary by CodeRabbit

Release Notes

• New Features
  • Added --batch-size CLI option to customize batch size for model analysis and VRAM footprint calculations (default: 1).
• Bug Fixes
  • Improved validation so --batch-size is rejected when set to 0 or a negative value.
• Tests
  • Added CLI tests to verify defaults, validation behavior, and that batch size is correctly applied to footprint calculations.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 2c0e93a2-2e37-4c86-92a6-91a6823c3462

📥 Commits

Reviewing files that changed from the base of the PR and between a260e36 and 7fce635.

📒 Files selected for processing (2)

• src/modelinfo/cli.py
• tests/test_cli.py

🚧 Files skipped from review as they are similar to previous changes (2)

• tests/test_cli.py
• src/modelinfo/cli.py

---

Walkthrough

The --batch-size CLI option (integer, default 1) is added to parse_args with a _positive_int validator that rejects values less than 1. The analyze_model function gains a matching batch_size: int = 1 parameter that it forwards to calculate_footprint. Both the single-file and multi-file paths in main pass args.batch_size into analyze_model. Tests verify default parsing, explicit parsing, error handling for invalid inputs, and end-to-end forwarding to calculate_footprint.

Changes

--batch-size flag: CLI, wiring, and tests

Layer / File(s)	Summary
CLI flag definition and function signature src/modelinfo/cli.py	_positive_int validator rejects batch sizes less than 1; parse_args registers --batch-size as an integer argument with default 1; analyze_model adds batch_size: int = 1 parameter.
Footprint calculation and main call site wiring src/modelinfo/cli.py	The calculate_footprint call receives batch_size=batch_size; both the multi-file comparison and single-file branches in main pass batch_size=args.batch_size into analyze_model.
Tests for --batch-size parsing and forwarding tests/test_cli.py	Imports modelinfo.cli as cli for monkeypatching; adds argument-parsing tests for default value (1), explicit integer acceptance, and zero/negative rejection; adds a behavior test that verifies analyze_model forwards batch_size and context_override to calculate_footprint and validates the returned kv_cache_bytes.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title 'feat(cli): add batch size option' clearly summarizes the main change - adding a CLI option for batch size control.
Description check	✅ Passed	The PR description covers the summary, references the linked issue (#2), and documents comprehensive testing performed including pytest, compilation, and linting checks.
Linked Issues check	✅ Passed	The PR fully addresses issue #2 by adding the --batch-size CLI option with default value 1, implementing input validation via _positive_int validator, and passing it through analyze_model to calculate_footprint with comprehensive test coverage.
Out of Scope Changes check	✅ Passed	All changes are directly related to adding and testing the --batch-size CLI option as specified in issue #2; no out-of-scope modifications are present.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

🧹 Nitpick comments (1)

src/modelinfo/cli.py (1)

55-60: ⚡ Quick win

Consider validating that batch size is positive.

The --batch-size argument accepts any integer, including zero or negative values. Since batch_size is multiplied into the KV cache calculation (kv_cache_bytes = 2 * num_layers * kv_dim * context_length * batch_size * 2), a value less than 1 would produce misleading or incorrect results (zero cache for batch_size=0, negative cache for negative values).

🛡️ Suggested validation

 parser.add_argument(
     "--batch-size",
     type=int,
     default=1,
     help="Batch size for dynamic KV cache footprint calculation.",
 )
+
+def _validate_batch_size(value: int) -> int:
+    if value < 1:
+        raise argparse.ArgumentTypeError("batch size must be at least 1")
+    return value

Then update the argument definition:

 parser.add_argument(
     "--batch-size",
-    type=int,
+    type=_validate_batch_size,
     default=1,
     help="Batch size for dynamic KV cache footprint calculation.",
 )

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/modelinfo/cli.py` around lines 55 - 60, The `--batch-size` argument in
the parser.add_argument call accepts any integer without validation, allowing
zero or negative values which produce incorrect KV cache calculations. Add
validation to the argument definition to ensure batch_size is a positive integer
(greater than 0). This can be done by adding a custom type function that
validates the input, or by using the choices parameter with a positive range
constraint, so that only valid positive batch sizes are accepted.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@src/modelinfo/cli.py`:
- Around line 55-60: The `--batch-size` argument in the parser.add_argument call
accepts any integer without validation, allowing zero or negative values which
produce incorrect KV cache calculations. Add validation to the argument
definition to ensure batch_size is a positive integer (greater than 0). This can
be done by adding a custom type function that validates the input, or by using
the choices parameter with a positive range constraint, so that only valid
positive batch sizes are accepted.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: e66852cc-71c5-4e08-9e90-b61c9a0d5c5d

📥 Commits

Reviewing files that changed from the base of the PR and between 33afb50 and a260e36.

📒 Files selected for processing (2)

• src/modelinfo/cli.py
• tests/test_cli.py

[pipe1os]

Thank you for your contribution!

your implementation is correct, I will proceed to merge this.