args.py 4.52 KB
Newer Older
1
2
3
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

4
"""CLI argument parsing for JSONL benchmark generation."""
5
6

import argparse
7
import sys
8
9
10
11
12
13
14
from pathlib import Path

DEFAULT_IMAGES_PER_REQUEST = 3
USER_TEXT_TOKENS = 300
COCO_ANNOTATIONS = Path(__file__).parent / "annotations" / "image_info_test2017.json"


15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
def _positive_int(value: str) -> int:
    iv = int(value)
    if iv <= 0:
        raise argparse.ArgumentTypeError(f"must be a positive integer, got {iv}")
    return iv


def _common_parser() -> argparse.ArgumentParser:
    """Args shared across all strategies."""
    p = argparse.ArgumentParser(add_help=False)
    p.add_argument(
        "-o",
        "--output",
        type=Path,
        default=None,
        help="Output .jsonl path (default: auto-generated from parameters)",
31
    )
32
33
    p.add_argument(
        "--user-text-tokens",
34
        type=int,
35
36
        default=USER_TEXT_TOKENS,
        help=f"Target user text tokens per request (default: {USER_TEXT_TOKENS})",
37
    )
38
39
    p.add_argument(
        "--seed",
40
41
        type=int,
        default=None,
42
        help="Random seed for reproducible generation (default: time-based)",
43
    )
44
45
46
47
48
49
50
51
    return p


def _image_parser() -> argparse.ArgumentParser:
    """Args for image generation (reusable for future video/audio parsers)."""
    p = argparse.ArgumentParser(add_help=False)
    p.add_argument(
        "--image-size",
52
        type=int,
53
54
55
56
        nargs=2,
        default=[512, 512],
        metavar=("WIDTH", "HEIGHT"),
        help="Size of generated PNG images in pixels (default: 512 512)",
57
    )
58
    p.add_argument(
59
60
61
62
63
        "--image-dir",
        type=Path,
        default=Path("/tmp/bench_images"),
        help="Directory to save generated PNG images (default: /tmp/bench_images)",
    )
64
    p.add_argument(
65
66
67
        "--image-mode",
        choices=["base64", "http"],
        default="base64",
68
        help="'base64' generates local PNGs (default); 'http' uses COCO URLs",
69
    )
70
    p.add_argument(
71
72
73
74
75
        "--coco-annotations",
        type=Path,
        default=COCO_ANNOTATIONS,
        help=f"Path to COCO image_info JSON for --image-mode http (default: {COCO_ANNOTATIONS})",
    )
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
    return p


def parse_args(description: str = "") -> argparse.Namespace:
    common = _common_parser()
    image = _image_parser()

    parser = argparse.ArgumentParser(
        description=description,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    sub = parser.add_subparsers(dest="strategy")

    # --- single-turn (default) ---
    st = sub.add_parser(
        "single-turn",
        parents=[common, image],
        help="Independent requests with random image sampling (default)",
    )
    st.add_argument(
        "-n",
        "--num-requests",
98
        type=int,
99
100
        default=500,
        help="Number of requests to generate (default: 500)",
101
    )
102
103
104
105
106
107
108
109
    st.add_argument(
        "--images-per-request",
        type=int,
        default=DEFAULT_IMAGES_PER_REQUEST,
        help=f"Number of images per request (default: {DEFAULT_IMAGES_PER_REQUEST})",
    )
    st.add_argument(
        "--images-pool",
110
111
        type=int,
        default=None,
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
        help="Unique images in pool. Smaller pool = more cross-request reuse. "
        "Default: num_requests * images_per_request (all unique).",
    )

    # --- sliding-window ---
    sw = sub.add_parser(
        "sliding-window",
        parents=[common, image],
        help="Causal sessions with sliding-window image overlap",
    )
    sw.add_argument(
        "--num-users",
        type=_positive_int,
        default=10,
        help="Number of concurrent user sessions (default: 10)",
127
    )
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
    sw.add_argument(
        "--turns-per-user",
        type=_positive_int,
        default=20,
        help="Number of requests per user (default: 20)",
    )
    sw.add_argument(
        "--window-size",
        type=_positive_int,
        default=5,
        help="Sliding window width — each turn sees this many images, "
        "with window_size-1 overlap between consecutive turns (default: 5)",
    )

    # Default to single-turn when no subcommand given, but let top-level
    # `-h`/`--help` flow through the main parser so users see both
    # subcommands and the module description.
    known_strategies = {"single-turn", "sliding-window"}
    argv = sys.argv[1:]
    help_requested = bool(argv) and argv[0] in {"-h", "--help"}
    if not help_requested and (not argv or argv[0] not in known_strategies):
        argv = ["single-turn", *argv]

    return parser.parse_args(argv)