From 551e566e4aebd1b8a52adc18e4873f9218bfe7fb Mon Sep 17 00:00:00 2001 From: Haris Khan Date: Sat, 8 Aug 2026 14:28:04 -0600 Subject: [PATCH] Add model benchmarking and cleanup examples --- examples/benchmark.py | 245 +++++++++++++++++++++++++++++++++++++ examples/modelDetection.py | 159 ++++++++++++++++++++++++ 2 files changed, 404 insertions(+) create mode 100644 examples/benchmark.py create mode 100644 examples/modelDetection.py diff --git a/examples/benchmark.py b/examples/benchmark.py new file mode 100644 index 00000000..526d1a9d --- /dev/null +++ b/examples/benchmark.py @@ -0,0 +1,245 @@ +"""Compare runtime performance across locally installed Ollama models.""" + +import argparse +import csv +import statistics +import sys +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, List, Optional, Tuple + +from ollama import GenerateResponse, ResponseError, generate +from ollama import list as list_models + +DEFAULT_PROMPT = 'Write a numbered list from 1 to 100, with each number followed by its English word.' +NANOSECONDS_PER_SECOND = 1_000_000_000 + + +@dataclass +class BenchmarkResult: + model: str + run: int + ttft_seconds: Optional[float] + wall_seconds: float + total_seconds: Optional[float] + load_seconds: Optional[float] + prompt_tokens: Optional[int] + prompt_tokens_per_second: Optional[float] + eval_tokens: Optional[int] + eval_tokens_per_second: Optional[float] + + +def non_negative_int(value): + parsed = int(value) + if parsed < 0: + raise argparse.ArgumentTypeError('must be zero or greater') + return parsed + + +def positive_int(value): + parsed = int(value) + if parsed < 1: + raise argparse.ArgumentTypeError('must be one or greater') + return parsed + + +def nanoseconds_to_seconds(value): + if value is None: + return None + return value / NANOSECONDS_PER_SECOND + + +def tokens_per_second(count, duration): + if count is None or duration is None or duration <= 0: + return None + return count * NANOSECONDS_PER_SECOND / duration + + +def benchmark_once(model, prompt, num_predict, keep_alive, run): + started = time.perf_counter() + first_token_at: Optional[float] = None + final_response: Optional[GenerateResponse] = None + + stream = generate( + model=model, + prompt=prompt, + stream=True, + keep_alive=keep_alive, + options={ + 'num_predict': num_predict, + 'seed': 42, + 'temperature': 0, + }, + ) + + for part in stream: + if first_token_at is None and (part.response or part.thinking): + first_token_at = time.perf_counter() + final_response = part + + finished = time.perf_counter() + if final_response is None: + raise RuntimeError('Ollama returned an empty response stream') + + return BenchmarkResult( + model=model, + run=run, + ttft_seconds=None if first_token_at is None else first_token_at - started, + wall_seconds=finished - started, + total_seconds=nanoseconds_to_seconds(final_response.total_duration), + load_seconds=nanoseconds_to_seconds(final_response.load_duration), + prompt_tokens=final_response.prompt_eval_count, + prompt_tokens_per_second=tokens_per_second(final_response.prompt_eval_count, final_response.prompt_eval_duration), + eval_tokens=final_response.eval_count, + eval_tokens_per_second=tokens_per_second(final_response.eval_count, final_response.eval_duration), + ) + + +def median(values): + present = [value for value in values if value is not None] + return statistics.median(present) if present else None + + +def summarize(model, results): + return { + 'model': model, + 'runs': len(results), + 'median_ttft_seconds': median(result.ttft_seconds for result in results), + 'median_wall_seconds': median(result.wall_seconds for result in results), + 'median_total_seconds': median(result.total_seconds for result in results), + 'median_load_seconds': median(result.load_seconds for result in results), + 'median_prompt_tokens_per_second': median(result.prompt_tokens_per_second for result in results), + 'median_eval_tokens_per_second': median(result.eval_tokens_per_second for result in results), + 'median_eval_tokens': median(None if result.eval_tokens is None else float(result.eval_tokens) for result in results), + } + + +def format_seconds(value): + if not isinstance(value, (float, int)): + return '-' + return f'{value:.3f}s' + + +def format_rate(value): + if not isinstance(value, (float, int)): + return '-' + return f'{value:.1f}' + + +def format_count(value): + if not isinstance(value, (float, int)): + return '-' + return f'{value:.0f}' + + +def print_table(summaries): + headers = ('Model', 'Runs', 'TTFT', 'Total', 'Load', 'Prompt tok/s', 'Gen tok/s', 'Gen tokens') + rows: List[Tuple[str, ...]] = [] + for summary in summaries: + rows.append( + ( + str(summary['model']), + str(summary['runs']), + format_seconds(summary['median_ttft_seconds']), + format_seconds(summary['median_total_seconds']), + format_seconds(summary['median_load_seconds']), + format_rate(summary['median_prompt_tokens_per_second']), + format_rate(summary['median_eval_tokens_per_second']), + format_count(summary['median_eval_tokens']), + ) + ) + + widths = [max(len(header), *(len(row[index]) for row in rows)) for index, header in enumerate(headers)] + print('\n' + ' '.join(header.ljust(widths[index]) for index, header in enumerate(headers))) + print(' '.join('-' * width for width in widths)) + for row in rows: + print(' '.join(value.ljust(widths[index]) for index, value in enumerate(row))) + + +def write_csv(path, summaries): + fieldnames = list(summaries[0]) + with path.open('w', newline='', encoding='utf-8') as csv_file: + writer = csv.DictWriter(csv_file, fieldnames=fieldnames) + writer.writeheader() + writer.writerows(summaries) + + +def parse_args(): + parser = argparse.ArgumentParser( + description='Compare runtime performance across locally installed Ollama models.', + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument('models', nargs='*', help='installed model names; omit to benchmark every installed model') + parser.add_argument('--prompt', default=DEFAULT_PROMPT, help='prompt sent to every model') + parser.add_argument('--runs', type=positive_int, default=3, help='measured runs per model') + parser.add_argument('--warmup', type=non_negative_int, default=1, help='unmeasured warm-up runs per model') + parser.add_argument('--num-predict', type=positive_int, default=128, help='maximum generated tokens per run') + parser.add_argument('--keep-alive', default='5m', help='how long Ollama keeps each model loaded') + parser.add_argument('--csv', type=Path, help='write median summary metrics to this CSV path') + return parser.parse_args() + + +def select_models(requested): + installed = sorted(model.model for model in list_models().models if model.model) + if not installed: + raise RuntimeError('No local models are installed. Run `ollama pull ` first.') + + if not requested: + return installed + + missing = sorted(set(requested) - set(installed)) + if missing: + raise RuntimeError(f'Models are not installed: {", ".join(missing)}') + return list(dict.fromkeys(requested)) + + +def main(): + args = parse_args() + try: + models = select_models(args.models) + except (ConnectionError, ResponseError, RuntimeError) as error: + print(f'Error: {error}', file=sys.stderr) + return 1 + + print(f'Benchmarking {len(models)} model(s) with {args.runs} measured run(s) each.') + print(f'Prompt: {args.prompt!r}') + print(f'Max generated tokens: {args.num_predict}; warm-up runs: {args.warmup}') + print('Keep hardware load and background activity stable for meaningful comparisons.') + + results_by_model: Dict[str, List[BenchmarkResult]] = {} + for model in models: + print(f'\n{model}') + try: + for warmup in range(1, args.warmup + 1): + print(f' warm-up {warmup}/{args.warmup}...', end='', flush=True) + benchmark_once(model, args.prompt, args.num_predict, args.keep_alive, run=0) + print(' done') + + model_results = [] + for run in range(1, args.runs + 1): + print(f' run {run}/{args.runs}...', end='', flush=True) + result = benchmark_once(model, args.prompt, args.num_predict, args.keep_alive, run=run) + model_results.append(result) + print(f' {format_rate(result.eval_tokens_per_second)} tokens/s') + results_by_model[model] = model_results + except (ConnectionError, ResponseError, RuntimeError) as error: + print(f' failed: {error}', file=sys.stderr) + + summaries = [summarize(model, results) for model, results in results_by_model.items()] + if not summaries: + print('No model completed the benchmark.', file=sys.stderr) + return 1 + + summaries.sort(key=lambda summary: float(summary['median_eval_tokens_per_second'] or 0), reverse=True) + print_table(summaries) + + if args.csv: + write_csv(args.csv, summaries) + print(f'\nWrote summary CSV to {args.csv}') + + return 0 + + +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/examples/modelDetection.py b/examples/modelDetection.py new file mode 100644 index 00000000..4ecfda32 --- /dev/null +++ b/examples/modelDetection.py @@ -0,0 +1,159 @@ +"""List installed Ollama models and optionally remove selected models.""" + +import sys +from dataclasses import dataclass + +from ollama import ResponseError +from ollama import delete as delete_model +from ollama import list as list_models + + +@dataclass(frozen=True) +class InstalledModel: + name: str + size: int + parameter_size: str + quantization: str + + +def format_size(size): + """Format a byte count using binary units.""" + value = float(size) + for unit in ('B', 'KiB', 'MiB', 'GiB', 'TiB'): + if value < 1024 or unit == 'TiB': + return f'{value:.1f} {unit}' + value /= 1024 + return f'{value:.1f} TiB' + + +def get_installed_models(): + """Return locally installed models reported by Ollama.""" + installed = [] + for model in list_models().models: + if not model.model: + continue + + details = model.details + installed.append( + InstalledModel( + name=model.model, + size=int(model.size) if model.size is not None else 0, + parameter_size=details.parameter_size if details and details.parameter_size else '-', + quantization=details.quantization_level if details and details.quantization_level else '-', + ) + ) + return sorted(installed, key=lambda model: model.name.lower()) + + +def print_models(models): + """Print installed models and their estimated storage usage.""" + headers = ('#', 'Model', 'Size', 'Parameters', 'Quantization') + rows = [(str(index), model.name, format_size(model.size), model.parameter_size, model.quantization) for index, model in enumerate(models, start=1)] + widths = [max(len(header), *(len(row[index]) for row in rows)) for index, header in enumerate(headers)] + + print('\nInstalled Ollama models:\n') + print(' '.join(header.ljust(widths[index]) for index, header in enumerate(headers))) + print(' '.join('-' * width for width in widths)) + for row in rows: + print(' '.join(value.ljust(widths[index]) for index, value in enumerate(row))) + + total_size = sum(model.size for model in models) + print(f'\nTotal reported size: {format_size(total_size)}') + + +def parse_selection(selection, models): + """Convert a comma-separated numeric selection into unique models.""" + normalized = selection.strip().lower() + if not normalized: + return [] + if normalized == 'all': + return list(models) + + selected = [] + seen = set() + for value in normalized.split(','): + value = value.strip() + if not value.isdigit(): + raise ValueError(f'Invalid selection: {value!r}') + + index = int(value) + if index < 1 or index > len(models): + raise ValueError(f'Model number {index} is outside the available range') + if index not in seen: + selected.append(models[index - 1]) + seen.add(index) + return selected + + +def ask_which_models(models): + """Prompt until the user chooses valid model numbers or cancels.""" + while True: + selection = input('\nEnter model numbers separated by commas, "all" to select everything,\nor press Enter to cancel: ') + try: + return parse_selection(selection, models) + except ValueError as error: + print(f'Error: {error}') + + +def confirm_removal(models): + """Require explicit confirmation before deleting models.""" + print('\nSelected for removal:') + for model in models: + print(f' - {model.name} ({format_size(model.size)})') + selected_size = sum(model.size for model in models) + print(f'\nSelected reported size: {format_size(selected_size)}') + print('Actual freed space may be smaller because Ollama models can share data layers.') + confirmation = input('Remove these models? Type "yes" to confirm: ') + return confirmation.strip().lower() == 'yes' + + +def remove_models(models): + """Remove each selected model and return the number of failures.""" + failures = 0 + for model in models: + try: + response = delete_model(model.name) + status = response.status or 'success' + print(f'Removed {model.name}: {status}') + except (ConnectionError, ResponseError) as error: + failures += 1 + print(f'Could not remove {model.name}: {error}', file=sys.stderr) + return failures + + +def main(): + try: + models = get_installed_models() + except (ConnectionError, ResponseError, KeyboardInterrupt) as error: + print(f'Could not connect to Ollama: {error}', file=sys.stderr) + print('Make sure Ollama is installed and running.', file=sys.stderr) + return 1 + + if not models: + print('No Ollama models are currently installed.') + return 0 + + print_models(models) + try: + selected = ask_which_models(models) + if not selected: + print('No models selected. Nothing was removed.') + return 0 + if not confirm_removal(selected): + print('Removal cancelled. Nothing was removed.') + return 0 + except (EOFError, KeyboardInterrupt): + print('\nRemoval cancelled. Nothing was removed.') + return 0 + + failures = remove_models(selected) + if failures: + print(f'Finished with {failures} removal failure(s).', file=sys.stderr) + return 1 + + print('\nSelected models were removed successfully.') + return 0 + + +if __name__ == '__main__': + raise SystemExit(main())