(
self,
source_ds: Dataset,
*,
per_process: Optional[_PerProcessFn] = None,
)
| 441 | """Tests Input module.""" |
| 442 | |
| 443 | def _input_config( |
| 444 | self, |
| 445 | source_ds: Dataset, |
| 446 | *, |
| 447 | per_process: Optional[_PerProcessFn] = None, |
| 448 | ): |
| 449 | def ds_fn() -> BuildDatasetFn: |
| 450 | def source(dispatch_config): |
| 451 | ds = source_ds |
| 452 | ds = shard_dataset(ds, dispatch_config=dispatch_config) |
| 453 | if callable(per_process): |
| 454 | ds = per_process(ds, dispatch_config=dispatch_config) |
| 455 | ds = prefetch_dataset( |
| 456 | maybe_to_iter_dataset(ds), |
| 457 | multiprocessing_options=grain.MultiprocessingOptions(num_workers=1), |
| 458 | ) |
| 459 | return ds |
| 460 | |
| 461 | return source |
| 462 | |
| 463 | cfg: Input.Config = Input.default_config().set(source=config_for_function(ds_fn)) |
| 464 | return cfg.set(name="test") |
| 465 | |
| 466 | @parameterized.parameters( |
| 467 | # A single process case. |
no test coverage detected