schrodinger.application.transforms.batching module

Shared batching transforms.

class schrodinger.application.transforms.batching.GlobalBatchConfig(*, batch_size: Optional[int] = None)

Bases: BaseModel

Configuration for GlobalBatch.

Parameters:

batch_size – Max elements per batch. When None, all elements are collected into a single batch.

batch_size: Optional[int]

model_config: ClassVar[ConfigDict] = {'frozen': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class schrodinger.application.transforms.batching.GlobalBatch(**kwargs)

Bases: PTransformWithConfig

Collect all elements globally and split into fixed-size batches.

beam.BatchElements batches per-worker and adapts batch sizes dynamically based on processing time, so the number and size of batches vary depending on the runner, worker count, and workload. This transform instead collects all elements into a single global list, then chunks deterministically – guaranteeing consistent batch sizes and count regardless of execution environment.

When batch_size is None, all elements are collected into a single batch (equivalent to beam.combiners.ToList()).

config_class

alias of GlobalBatchConfig