@@ -247,6 +247,7 @@ async def resolve_request_loader(
247247 data_num_workers : int | None ,
248248 random_seed : int ,
249249 console : Console | None = None ,
250+ max_requests : int | None = None ,
250251 ** dataloader_kwargs : dict [str , Any ] | None ,
251252) -> DataLoader [GenerationRequest ]:
252253 """
@@ -270,6 +271,7 @@ async def resolve_request_loader(
270271 :param data_num_workers: Number of worker processes for data loading
271272 :param random_seed: Seed for reproducible random operations
272273 :param console: Console instance for progress reporting, or None
274+ :param max_requests: If set, first data source loads at most this many rows.
273275 :param dataloader_kwargs: Additional arguments passed to DataLoader initialization
274276 :return: Configured DataLoader instance for GenerationRequest objects
275277 :raises ValueError: If request formatter type is not registered in
@@ -306,6 +308,17 @@ async def resolve_request_loader(
306308 data_finalizer ,
307309 )
308310
311+ # When max_requests is set, limit the first data source to that many rows at load
312+ if max_requests is not None and data :
313+ if max_requests < 1 :
314+ raise ValueError (
315+ "max_requests must be >= 1 when set for data truncation, "
316+ f"got { max_requests } "
317+ )
318+ data_args = list (data_args ) if data_args else [{} for _ in data ]
319+ if len (data_args ) >= 1 :
320+ data_args [0 ] = {** data_args [0 ], "max_rows" : max_requests }
321+
309322 request_loader : DataLoader [GenerationRequest ] = DataLoader (
310323 data = data ,
311324 data_args = data_args ,
@@ -352,6 +365,7 @@ async def resolve_profile(
352365 max_global_error_rate : float | None ,
353366 over_saturation : dict [str , Any ] | None = None ,
354367 console : Console | None = None ,
368+ data : list [Any ] | None = None ,
355369) -> Profile :
356370 """
357371 Resolve and configure a benchmark profile with rate and constraint settings.
@@ -373,6 +387,7 @@ async def resolve_profile(
373387 :param max_global_error_rate: Maximum global error rate threshold before stopping
374388 :param over_saturation: Over-saturation detection configuration (dict)
375389 :param console: Console instance for progress reporting, or None
390+ :param data: Optional list of data sources.
376391 :return: Configured Profile instance ready for benchmarking
377392 :raises ValueError: If constraints are provided with a pre-configured Profile
378393 """
@@ -400,6 +415,7 @@ async def resolve_profile(
400415 random_seed = random_seed ,
401416 rampup_duration = rampup ,
402417 constraints = {** constraints },
418+ data = data ,
403419 )
404420 elif constraints :
405421 raise ValueError (
@@ -488,24 +504,58 @@ async def benchmark_generative_text(
488504 processor = await resolve_processor (
489505 processor = args .processor , model = model , console = console
490506 )
491- request_loader = await resolve_request_loader (
492- data = args .data ,
493- model = model ,
494- data_args = args .data_args ,
495- data_samples = args .data_samples ,
496- processor = processor ,
497- processor_args = args .processor_args ,
498- data_column_mapper = args .data_column_mapper ,
499- data_preprocessors = args .data_preprocessors ,
500- data_preprocessors_kwargs = args .data_preprocessors_kwargs ,
501- data_finalizer = args .data_finalizer ,
502- data_collator = args .data_collator ,
503- data_sampler = args .data_sampler ,
504- data_num_workers = args .data_num_workers ,
505- random_seed = args .random_seed ,
506- console = console ,
507- ** (args .dataloader_kwargs or {}),
508- )
507+
508+ # Build common kwargs for resolve_profile and resolve_request_loader
509+ profile_kwargs = {
510+ "profile" : args .profile ,
511+ "rate" : args .rate ,
512+ "random_seed" : args .random_seed ,
513+ "rampup" : args .rampup ,
514+ "constraints" : constraints ,
515+ "max_seconds" : args .max_seconds ,
516+ "max_requests" : args .max_requests ,
517+ "max_errors" : args .max_errors ,
518+ "max_error_rate" : args .max_error_rate ,
519+ "max_global_error_rate" : args .max_global_error_rate ,
520+ "over_saturation" : args .over_saturation ,
521+ "console" : console ,
522+ }
523+ loader_kwargs = {
524+ "data" : args .data ,
525+ "model" : model ,
526+ "data_args" : args .data_args ,
527+ "data_samples" : args .data_samples ,
528+ "processor" : processor ,
529+ "processor_args" : args .processor_args ,
530+ "data_column_mapper" : args .data_column_mapper ,
531+ "data_preprocessors" : args .data_preprocessors ,
532+ "data_preprocessors_kwargs" : args .data_preprocessors_kwargs ,
533+ "data_finalizer" : args .data_finalizer ,
534+ "data_collator" : args .data_collator ,
535+ "data_sampler" : args .data_sampler ,
536+ "data_num_workers" : args .data_num_workers ,
537+ "random_seed" : args .random_seed ,
538+ "console" : console ,
539+ }
540+
541+ # For replay profile: resolve profile first to apply max_seconds filtering,
542+ # then use the filtered count for the data loader. This ensures the data
543+ # loader and scheduler both work with the same filtered request count.
544+ if args .profile == "replay" :
545+ profile = await resolve_profile (** profile_kwargs , data = args .data ) # type: ignore[arg-type]
546+ effective_max_requests = (
547+ profile .constraints .get ("max_requests" )
548+ if profile .constraints
549+ else args .max_requests
550+ )
551+ request_loader = await resolve_request_loader (
552+ ** loader_kwargs , max_requests = effective_max_requests
553+ ) # type: ignore[arg-type]
554+ else :
555+ request_loader = await resolve_request_loader (
556+ ** loader_kwargs , max_requests = args .max_requests
557+ ) # type: ignore[arg-type]
558+ profile = await resolve_profile (** profile_kwargs , data = None ) # type: ignore[arg-type]
509559
510560 warmup = TransientPhaseConfig .create_from_value (args .warmup )
511561 cooldown = TransientPhaseConfig .create_from_value (args .cooldown )
@@ -521,21 +571,6 @@ async def benchmark_generative_text(
521571 ),
522572 status = "success" ,
523573 )
524-
525- profile = await resolve_profile (
526- profile = args .profile ,
527- rate = args .rate ,
528- random_seed = args .random_seed ,
529- rampup = args .rampup ,
530- constraints = constraints ,
531- max_seconds = args .max_seconds ,
532- max_requests = args .max_requests ,
533- max_errors = args .max_errors ,
534- max_error_rate = args .max_error_rate ,
535- max_global_error_rate = args .max_global_error_rate ,
536- over_saturation = args .over_saturation ,
537- console = console ,
538- )
539574 output_formats = await resolve_output_formats (
540575 outputs = args .outputs , output_dir = args .output_dir , console = console
541576 )
0 commit comments