From 7f4a4db995da00cc85bcac86b73327f7a5c2cb55 Mon Sep 17 00:00:00 2001 From: Sean O'Brien <60306702+stobrien89@users.noreply.github.com> Date: Sun, 25 Oct 2026 18:57:01 -0400 Subject: [PATCH 1/3] partition schemas and model stubs --- designs/codegen/index.md | 1 + designs/codegen/partitioning.md | 163 +++++++++++ designs/codegen/types.md | 28 +- .../src/smithy_python/generation.py | 9 + .../smithy-python/src/smithy_python/models.py | 124 ++++++++- .../src/smithy_python/schemas.py | 260 ++++++++++++++++-- .../tests/unit/test_generation.py | 9 + .../smithy-python/tests/unit/test_models.py | 85 +++++- .../smithy-python/tests/unit/test_schemas.py | 167 ++++++++++- 9 files changed, 792 insertions(+), 54 deletions(-) create mode 100644 designs/codegen/partitioning.md diff --git a/designs/codegen/index.md b/designs/codegen/index.md index 5455a37b9..81d968097 100644 --- a/designs/codegen/index.md +++ b/designs/codegen/index.md @@ -65,3 +65,4 @@ behavior of generated packages. * [Python Symbols](symbols.md) * [Python Writer](writer.md) * [Standalone Types Generation](types.md) +* [Code Generation Partitioning](partitioning.md) diff --git a/designs/codegen/partitioning.md b/designs/codegen/partitioning.md new file mode 100644 index 000000000..e4f0c8266 --- /dev/null +++ b/designs/codegen/partitioning.md @@ -0,0 +1,163 @@ +# Code Generation Partitioning + +## Abstract + +Large standalone types packages split compact schemas and model stubs across +private modules. Stable facade modules preserve generated imports while +reducing generator memory and static analysis time. + +## Motivation + +An EC2 types package contains 4,065 selected shapes. Emitting its schemas and +model declarations as two large files makes each file the unit of work for +rendering, staging, parsing, and static analysis. + +`FileManifest` releases each rendered file after staging it, so the largest file +sets the lower bound for source memory during publication. Pyright analyzes the +model stub as one large declaration file. Lazy schema hydration and model +construction reduce runtime work but do not reduce these file-level costs. + +Private partitions bound the work associated with one generated file. Stable +facades preserve generated imports and lazy runtime behavior while tools process +the smaller partition files. + +## Settings + +`TypesSettings` owns both partition targets: + +```python +@dataclass(frozen=True, slots=True) +class TypesSettings(ArtifactSettings): + selector: Selector = field(default_factory=Selector.identity) + generate_inputs_and_outputs: bool = False + schema_partition_size: int | None = 250 + stub_partition_size: int | None = 250 +``` + +The CLI uses these defaults. Python callers may set either target to a positive +integer or `None`. A value of `None` emits that surface as one file. + +Partitioning starts when the number of planned items exceeds its target. A +package with 250 schema shapes and 250 model declarations keeps the single-file +layout. + +## Output Layout + +Partitioned packages retain the stable module paths used by unpartitioned +packages. + +| Surface | Facade | Partitions | +| --- | --- | --- | +| Compact schemas | `._private.schemas` | `._private.schemas_N` | +| Model stubs | `.models.pyi` | `._private.models_N.pyi` | +| Runtime models | `.models` | Single module | + +### Schema File Layout + +```text +src// +├── models.py +└── _private/ + ├── __init__.py + ├── schemas.py + ├── schemas_0.py + ├── schemas_1.py + └── schemas_N.py +``` + +`models.py` imports the `schemas.py` facade. The facade maps schema constants to +partition indexes and loads the matching `schemas_N.py` module on first access. +Each partition contains compact schema tuples. + +### Stub File Layout + +```text +src// +├── models.py +├── models.pyi +└── _private/ + ├── __init__.py + ├── models_0.pyi + ├── models_1.pyi + └── models_N.pyi +``` + +`models.pyi` is the public typing facade. It imports declarations from each +private stub partition. `models.py` remains the runtime module. + +`FileManifest` registers the facade and partition files before publication. +Its ownership rules remove stale, unmodified partition files after a model +shrinks or a target changes. + +Files under `._private` remain private implementation details. + +## Schema Partitions + +`SchemaPlan` provides schema shapes in deterministic emission order and groups +each recursive strongly connected component. A recursive group stays in one +partition so its schemas can use direct local references. + +The schema partitioner: + +1. Packs recursive groups in plan order using the configured shape target. +2. Uses the resulting partition count for the final layout. +3. Estimates the emitted source size of each group. +4. Selects contiguous boundaries that balance estimated source size. + +The estimate includes shape and member identifiers, schema references, and +trait values. It uses plan metadata instead of rendering schemas a second time. +The target determines the partition count. Final partitions may exceed the +target to balance source size, and recursive groups remain indivisible. + +The facade maps each schema constant to a partition index. Its `__getattr__` +imports the selected partition, reads the constant, and caches the result in +the facade module. `__dir__` includes every planned schema constant. + +Partition modules route cross-partition references through the facade. This +keeps dependencies lazy and avoids direct import cycles between partitions. +Hydration remains lazy through `smithy_core._schema_compact`. + +## Stub Partitions + +`ModelPlan` provides model declarations in deterministic order. The stub +partitioner writes contiguous groups up to the configured target. + +The public `models.pyi` facade imports every declaration from its private stub +partition and defines the same `__all__` as `models.py`. Private stubs reference +generated declarations through the public `models` module. This keeps their +type references independent of private partition boundaries. + +Runtime `models.py` remains one lazy registry module. Runtime model construction +and class identity stay unchanged. + +## Compatibility + +Partitioning preserves: + +* Public model module paths and runtime schema lookups. +* Model names, constructor signatures, and type annotations. +* Public model `__all__`, explicit imports, star imports, and `dir()` results. +* Lazy schema hydration and model construction. +* Schema identity, recursive references, and serde behavior. + +Partition indexes follow generation-plan order. Model changes can move +declarations between private files while the public facades remain stable. + +## Validation + +Tests cover packages below and above each target, recursive schema groups, +cross-partition references, facade caching, strict Pyright, Ruff, and +deterministic output. + +The final x86 EC2 benchmark compared unpartitioned output with 250-item schema +and stub partitions: + +| Metric | Unpartitioned | Partitioned | Change | +| --- | ---: | ---: | ---: | +| Generation time | 3.557 s | 3.555 s | -0.1% | +| Peak RSS | 264.5 MiB | 122.4 MiB | -53.7% | +| Peak Python allocation | 222.6 MiB | 84.7 MiB | -61.9% | +| Generated bytes | 4.78 MiB | 5.33 MiB | +11.6% | +| Model import | 50.6 ms | 40.4 ms | -20.2% | +| First model access | 12.3 ms | 22.3 ms | +10.0 ms | +| Strict Pyright | 81.98 s | 15.25 s | -81.4% | diff --git a/designs/codegen/types.md b/designs/codegen/types.md index 58e07fa24..964c4b6f8 100644 --- a/designs/codegen/types.md +++ b/designs/codegen/types.md @@ -100,11 +100,15 @@ class ArtifactSettings: class TypesSettings(ArtifactSettings): selector: Selector = field(default_factory=Selector.identity) generate_inputs_and_outputs: bool = False + schema_partition_size: int | None = 250 + stub_partition_size: int | None = 250 ``` The resolved service remains on `Selection`. It is not duplicated in settings. The CLI subcommand identifies the artifact, so settings do not duplicate that value. Artifact-specific settings extend the shared frozen dataclass. +Partition targets and their output layout are defined in +[Code Generation Partitioning](partitioning.md). The `types` command accepts: @@ -416,26 +420,12 @@ Package setup renders the dependency set collected during generation. ## Large Models -Schema and stub partitioning must not change public imports. The output API and -generation plan allow private content to move behind stable facade modules. +Large standalone types packages partition compact schemas and model stubs +behind stable facade modules. Schema partitions use source-size weighting and +keep recursive components together. Stub partitions use declaration counts. -smithy-java partitions schema classes after a model reaches 500 schema fields -and then targets about 50 KB of generated schema content per class. smithy-php -uses a centralized lazy schema registry. Python can combine those approaches: - -* Keep `._private.schemas` as the stable facade. -* Place compact schema tuples in private partition modules. -* Keep every recursive strongly connected component in one partition. -* Route schema constant access through generated indexes. -* Partition large type stubs behind the public `models.pyi` facade. - -Python partitioning thresholds must be based on Python measurements. Relevant inputs include writer render time, peak generator memory, generated import time, and type -checker time. - -Partitioning bounds the syntax tree analyzed by `PythonWriter` and prevents a -single EC2-sized source file from dominating codegen memory. Measurements -establish the partitioning thresholds. Facade modules preserve generated -imports. +[Code Generation Partitioning](partitioning.md) defines the output layout, +thresholds, and compatibility requirements. ## Performance diff --git a/packages/smithy-python/src/smithy_python/generation.py b/packages/smithy-python/src/smithy_python/generation.py index ce75e1e19..b82369040 100644 --- a/packages/smithy-python/src/smithy_python/generation.py +++ b/packages/smithy-python/src/smithy_python/generation.py @@ -44,6 +44,15 @@ class TypesSettings(ArtifactSettings): selector: _selection.Selector = field(default_factory=_selection.Selector.identity) generate_inputs_and_outputs: bool = False + schema_partition_size: int | None = 250 + stub_partition_size: int | None = 250 + + def __post_init__(self) -> None: + if any( + size is not None and size < 1 + for size in (self.schema_partition_size, self.stub_partition_size) + ): + raise ValueError("Partition sizes must be positive") @dataclass(frozen=True, slots=True) diff --git a/packages/smithy-python/src/smithy_python/models.py b/packages/smithy-python/src/smithy_python/models.py index d6ae2cd48..1c346472d 100644 --- a/packages/smithy-python/src/smithy_python/models.py +++ b/packages/smithy-python/src/smithy_python/models.py @@ -36,12 +36,31 @@ def generate(self, context: GenerationContext[TypesSettings]) -> None: ), context.plan.models, ) - _StubEmitter( + facade = _StubEmitter( context.files.python_module( f"src/{directory}/models.pyi", f"{package}.models" ), context, - ).write() + ) + shapes = context.plan.models.shapes + size = context.settings.stub_partition_size + if size is None or len(shapes) <= size: + facade.write() + return + + partitions = tuple( + shapes[index : index + size] for index in range(0, len(shapes), size) + ) + facade.write_facade(partitions) + for index, shapes in enumerate(partitions): + _StubEmitter( + context.files.python_module( + f"src/{directory}/_private/models_{index}.pyi", + f"{package}._private.models_{index}", + ), + context, + external_bases=True, + ).write_partition(shapes) def _core_reference(name: str, module: str) -> _symbols.TypeReference: @@ -165,12 +184,25 @@ def __init__( self, writer: PythonWriter, context: GenerationContext[TypesSettings], + *, + external_bases: bool = False, ) -> None: self.writer = writer self.model = context.model self.plan = context.plan.models self.symbols = context.symbols self.module = f"{context.settings.package}.models" + self.external_bases = external_bases + self.serde_shape: str | _symbols.TypeReference = ( + _symbols.TypeReference("_SerdeShape", "..models") + if external_bases + else "_SerdeShape" + ) + self.service_error: str | _symbols.TypeReference = ( + _symbols.TypeReference("ServiceError", "..models") + if external_bases + else "ServiceError" + ) operations = tuple(context.model.iter_shapes(_model.OperationShape)) self.inputs = {operation.input for operation in operations} | { shape.id for shape in self.plan.shapes if shape.has_trait("input") @@ -180,6 +212,30 @@ def __init__( } def write(self) -> None: + """Write the complete unpartitioned public stub.""" + self._write_preamble() + self._write_shapes(self.plan.shapes) + self._write_unit() + + def write_facade(self, partitions: tuple[tuple[_model.Shape, ...], ...]) -> None: + """Write the public stub facade and re-export private declarations.""" + writer = self.writer + for index, shapes in enumerate(partitions): + writer.line(f"from ._private.models_{index} import (") + with writer.indent(): + for shape in shapes: + for name in _stub_shape_names(self.plan, shape): + writer.line(name, ",") + writer.line(")") + writer.line() + self._write_preamble() + self._write_unit() + + def write_partition(self, shapes: tuple[_model.Shape, ...]) -> None: + """Write declarations assigned to one private stub partition.""" + self._write_shapes(shapes) + + def _write_preamble(self) -> None: writer = self.writer writer.line("__all__ = [") with writer.indent(): @@ -212,7 +268,9 @@ def write(self) -> None: with writer.indent(): _write_docstring(writer, _SERVICE_ERROR_DOC) - for shape in self.plan.shapes: + def _write_shapes(self, shapes: Collection[_model.Shape]) -> None: + writer = self.writer + for shape in shapes: writer.line() match shape.type: case _model.ShapeType.STRUCTURE: @@ -223,10 +281,13 @@ def write(self) -> None: self._write_enum(shape) case _: pass + + def _write_unit(self) -> None: + writer = self.writer if self.plan.includes_unit: writer.line() writer.line("@", _DATACLASS, "(kw_only=True)") - writer.line("class Unit(_SerdeShape):") + writer.line("class Unit(", self.serde_shape, "):") with writer.indent(): writer.line("pass") @@ -235,10 +296,17 @@ def _write_structure(self, shape: _model.Shape) -> None: name = self.plan.declaration_names[shape.id] error = shape.has_trait("error") writer.line("@", _DATACLASS, "(kw_only=True)") + bases: tuple[str | _symbols.TypeReference, ...] = ( + (self.service_error, ", ", self.serde_shape) + if error + else (self.serde_shape,) + ) writer.line( "class ", name, - "(ServiceError, _SerdeShape):" if error else "(_SerdeShape):", + "(", + *bases, + "):", ) with writer.indent(): members = tuple( @@ -255,7 +323,7 @@ def _write_structure(self, shape: _model.Shape) -> None: if members or carries_metadata: writer.line() for member in members: - reference = self.symbols.type_reference(member.target) + reference = self._type_reference(member.target) has_default = member.has_trait("default") or self.model.get_target( member ).has_trait("default") @@ -283,20 +351,20 @@ def _write_union(self, shape: _model.Shape) -> None: variants: list[_symbols.TypeReference] = [] for member in shape.members.values(): name = self.plan.union_variant_names[member.id] - variants.append(_symbols.TypeReference(name, self.module)) + variants.append(self._public_reference(name)) writer.line("@", _DATACLASS) - writer.line("class ", name, "(_SerdeShape):") + writer.line("class ", name, "(", self.serde_shape, "):") with writer.indent(): if docs := _documentation(member): _write_docstring(writer, docs) writer.line() - writer.line("value: ", self.symbols.type_reference(member.target)) + writer.line("value: ", self._type_reference(member.target)) writer.line() unknown = self.plan.union_unknown_names[shape.id] - variants.append(_symbols.TypeReference(unknown, self.module)) + variants.append(self._public_reference(unknown)) writer.line("@", _DATACLASS) - writer.line("class ", unknown, "(_SerdeShape):") + writer.line("class ", unknown, "(", self.serde_shape, "):") with writer.indent(): _write_docstring(writer, _UNKNOWN_UNION_DOC) writer.line() @@ -314,7 +382,7 @@ def _write_union(self, shape: _model.Shape) -> None: "def deserialize(self, deserializer: ", _SHAPE_DESERIALIZER, ") -> ", - _symbols.TypeReference(name, self.module), + self._public_reference(name), ": ...", ) @@ -341,12 +409,44 @@ def _write_enum(self, shape: _model.Shape) -> None: if not docs and not shape.members: writer.line("pass") + def _type_reference(self, shape_id: _model.ShapeId) -> _symbols.TypeReference: + return self._partition_reference(self.symbols.type_reference(shape_id)) + + def _public_reference(self, name: str) -> _symbols.TypeReference: + return _symbols.TypeReference( + name, "..models" if self.external_bases else self.module + ) + + def _partition_reference( + self, reference: _symbols.TypeReference + ) -> _symbols.TypeReference: + arguments = tuple(map(self._partition_reference, reference.arguments)) + module = ( + "..models" + if self.external_bases and reference.module == self.module + else reference.module + ) + if module == reference.module and arguments == reference.arguments: + return reference + return replace(reference, module=module, arguments=arguments) + def _model_exports(plan: ModelPlan) -> tuple[str, ...]: exports = ("ServiceError", *plan.registry_index) return (*exports, "Unit") if plan.includes_unit else exports +def _stub_shape_names(plan: ModelPlan, shape: _model.Shape) -> tuple[str, ...]: + declaration = plan.declaration_names[shape.id] + if shape.type is not _model.ShapeType.UNION: + return (declaration,) + return ( + *(plan.union_variant_names[member.id] for member in shape.members.values()), + plan.union_unknown_names[shape.id], + declaration, + ) + + def _write_mapping(writer: PythonWriter, name: str, values: Mapping[str, str]) -> None: if not values: writer.line(name, " = {}") diff --git a/packages/smithy-python/src/smithy_python/schemas.py b/packages/smithy-python/src/smithy_python/schemas.py index c47084e45..cf0764641 100644 --- a/packages/smithy-python/src/smithy_python/schemas.py +++ b/packages/smithy-python/src/smithy_python/schemas.py @@ -27,6 +27,7 @@ class SchemaPlan: """Derived semantic data used to emit compact schemas.""" shapes: tuple[_model.Shape, ...] + partition_groups: tuple[tuple[_model.Shape, ...], ...] names: Mapping[_model.ShapeId, str] strings: Mapping[str, str] traits: Mapping[_model.ShapeId, _TraitPlan] @@ -39,7 +40,7 @@ def from_shapes( symbols: _symbols.SymbolProvider, ) -> SchemaPlan: """Plan schema order, filtered traits, names, and interned strings.""" - emitted = _schema_order(shapes) + emitted, partition_groups = _schema_order(shapes) included = _included_traits(model) inputs = { operation.input for operation in model.iter_shapes(_model.OperationShape) @@ -69,6 +70,7 @@ def from_shapes( return cls( shapes=emitted, + partition_groups=partition_groups, names=MappingProxyType( {shape.id: symbols.schema_name(shape.id) for shape in emitted} ), @@ -83,18 +85,111 @@ class SchemaGenerator: def generate(self, context: GenerationContext[TypesSettings]) -> None: """Register ``._private.schemas`` with the file manifest.""" package = context.settings.package - path = f"src/{package.replace('.', '/')}/_private/schemas.py" - writer = context.files.python_module(path, f"{package}._private.schemas") + directory = package.replace(".", "/") plan = context.plan.schemas - for value, name in plan.strings.items(): + size = context.settings.schema_partition_size + if size is None or len(plan.shapes) <= size: + writer = context.files.python_module( + f"src/{directory}/_private/schemas.py", + f"{package}._private.schemas", + ) + _write_schema_module(writer, plan, plan.shapes) + return + + partitions = _schema_partitions(plan, size) + _write_schema_facade( + context.files.python_module( + f"src/{directory}/_private/schemas.py", + f"{package}._private.schemas", + ), + plan, + partitions, + ) + for index, shapes in enumerate(partitions): + _write_schema_module( + context.files.python_module( + f"src/{directory}/_private/schemas_{index}.py", + f"{package}._private.schemas_{index}", + ), + plan, + shapes, + facade="schemas", + ) + + +_IMPORT_MODULE: Final = _symbols.TypeReference("import_module", "importlib") + + +def _write_schema_facade( + writer: PythonWriter, + plan: SchemaPlan, + partitions: tuple[tuple[_model.Shape, ...], ...], +) -> None: + writer.line("_PARTITION_NAMES: tuple[tuple[str, ...], ...] = (") + with writer.indent(): + for shapes in partitions: + writer.line("(") + with writer.indent(): + for shape in shapes: + writer.line(repr(plan.names[shape.id]), ",") + writer.line("),") + writer.line(")") + writer.line("_PARTITIONS = {") + with writer.indent(): + writer.line("name: index") + writer.line("for index, names in enumerate(_PARTITION_NAMES)") + writer.line("for name in names") + writer.line("}") + writer.line("del _PARTITION_NAMES") + writer.line("def __getattr__(name: str):") + with writer.indent(): + writer.line("try:") + with writer.indent(): + writer.line("partition = _PARTITIONS[name]") + writer.line("except KeyError as error:") + with writer.indent(): + writer.line("raise AttributeError(name) from error") + writer.line( + "module = ", + _IMPORT_MODULE, + '(f".schemas_{partition}", __package__)', + ) + writer.line("value = getattr(module, name)") + writer.line("globals()[name] = value") + writer.line("return value") + writer.line() + writer.line() + writer.line("def __dir__() -> list[str]:") + with writer.indent(): + writer.line("return sorted({*globals(), *_PARTITIONS})") + + +def _write_schema_module( + writer: PythonWriter, + plan: SchemaPlan, + shapes: tuple[_model.Shape, ...], + *, + facade: str | None = None, +) -> None: + local = frozenset(shape.id for shape in shapes) + if facade is not None and any( + member.target in plan.names and member.target not in local + for shape in shapes + for member in shape.members.values() + ): + writer.line("from . import ", facade) + writer.line() + strings = _schema_strings(plan, shapes) + for value, name in plan.strings.items(): + if value in strings: writer.line(name, " = ", repr(value)) - if plan.strings: - writer.line() + if strings: + writer.line() - written: set[_model.ShapeId] = set() - for shape in plan.shapes: - _write_shape(writer, plan, shape, written) - written.add(shape.id) + written: set[_model.ShapeId] = set() + for shape in shapes: + _write_shape(writer, plan, shape, written, local=local, facade=facade) + written.add(shape.id) # This contract is intentionally local to the generator. It must match @@ -264,6 +359,9 @@ def _write_shape( plan: SchemaPlan, shape: _model.Shape, written: set[_model.ShapeId], + *, + local: Collection[_model.ShapeId], + facade: str | None, ) -> None: writer.line(plan.names[shape.id], ": ", _STATIC_SCHEMA, " = (") with writer.indent(): @@ -284,7 +382,13 @@ def _write_shape( for member in shape.members.values(): writer.line( "(", - _schema_ref(plan, member.target, written), + *_schema_ref( + plan, + member.target, + written, + local=local, + facade=facade, + ), ", ", *_trait_parts(plan, plan.traits[member.id]), "),", @@ -406,23 +510,36 @@ def _node_parts( def _schema_ref( - plan: SchemaPlan, target: _model.ShapeId, written: set[_model.ShapeId] -) -> str: + plan: SchemaPlan, + target: _model.ShapeId, + written: set[_model.ShapeId], + *, + local: Collection[_model.ShapeId], + facade: str | None, +) -> tuple[_SourcePart, ...]: if (prelude := _PRELUDE_INTS.get(target)) is not None: - return str(prelude) + return (str(prelude),) try: name = plan.names[target] except KeyError as error: raise _exceptions.CodegenError( f"Schema target {target} was not selected" ) from error - return name if target in written else f"lambda: {name}" + if target in local: + return (name if target in written else f"lambda: {name}",) + if facade is None: + raise _exceptions.CodegenError(f"Schema target {target} is not in this module") + return ("lambda: ", facade, ".", name) def _schema_order( shapes: tuple[_model.Shape, ...], -) -> tuple[_model.Shape, ...]: +) -> tuple[ + tuple[_model.Shape, ...], + tuple[tuple[_model.Shape, ...], ...], +]: by_id = {shape.id: shape for shape in shapes} + positions = {shape.id: index for index, shape in enumerate(shapes)} graph = { shape.id: tuple( member.target for member in shape.members.values() if member.target in by_id @@ -457,6 +574,7 @@ def _schema_order( pending.append((target, 0)) recursive: set[_model.ShapeId] = set() + recursive_components: dict[_model.ShapeId, tuple[_model.ShapeId, ...]] = {} visited.clear() for root in reversed(finished): if root in visited: @@ -473,11 +591,97 @@ def _schema_order( component_pending.append(target) if len(component) > 1 or root in graph[root]: recursive.update(component) + ordered_component = tuple(sorted(component, key=positions.__getitem__)) + for shape_id in component: + recursive_components[shape_id] = ordered_component - return ( + emitted = ( *(by_id[shape_id] for shape_id in finished if shape_id not in recursive), *(shape for shape in shapes if shape.id in recursive), ) + groups: list[tuple[_model.Shape, ...]] = [] + grouped: set[_model.ShapeId] = set() + for shape in emitted: + if shape.id in grouped: + continue + group_ids = recursive_components.get(shape.id, (shape.id,)) + groups.append(tuple(by_id[shape_id] for shape_id in group_ids)) + grouped.update(group_ids) + return emitted, tuple(groups) + + +def _schema_partitions( + plan: SchemaPlan, + size: int, +) -> tuple[tuple[_model.Shape, ...], ...]: + groups = plan.partition_groups + partitions: list[tuple[_model.Shape, ...]] = [] + current: list[_model.Shape] = [] + for group in groups: + if current and len(current) + len(group) > size: + partitions.append(tuple(current)) + current.clear() + current.extend(group) + if current: + partitions.append(tuple(current)) + partition_count = len(partitions) + if partition_count < 2: + return tuple(partitions) + + weights = tuple(_schema_group_weight(plan, group) for group in groups) + partitions.clear() + start = 0 + remaining_weight = sum(weights) + for remaining_parts in range(partition_count, 1, -1): + target = remaining_weight / remaining_parts + end = start + weight = 0 + last_end = len(groups) - remaining_parts + 1 + while end < last_end: + next_weight = weight + weights[end] + if weight and abs(target - weight) <= abs(target - next_weight): + break + weight = next_weight + end += 1 + partitions.append( + tuple(shape for group in groups[start:end] for shape in group) + ) + start = end + remaining_weight -= weight + partitions.append(tuple(shape for group in groups[start:] for shape in group)) + return tuple(partitions) + + +def _schema_group_weight(plan: SchemaPlan, group: tuple[_model.Shape, ...]) -> int: + declarations = sum( + len(plan.strings[value]) + len(repr(value)) + 4 + for value in _schema_strings(plan, group) + ) + return declarations + sum(_schema_shape_weight(plan, shape) for shape in group) + + +def _schema_shape_weight(plan: SchemaPlan, shape: _model.Shape) -> int: + """Estimate compact-schema source bytes without rendering shapes twice.""" + return ( + 60 + + len(shape.id.namespace) + + 2 * len(shape.id.name) + + _source_size(_trait_parts(plan, plan.traits[shape.id])) + + sum( + 40 + + 2 * len(member.name) + + len(plan.names.get(member.target, member.target.name)) + + _source_size(_trait_parts(plan, plan.traits[member.id])) + for member in shape.members.values() + ) + ) + + +def _source_size(parts: Collection[_SourcePart]) -> int: + return sum( + len(part.name if isinstance(part, _symbols.TypeReference) else part) + for part in parts + ) def _intern_traits(strings: _StringStore, traits: _TraitPlan) -> None: @@ -489,6 +693,28 @@ def _intern_traits(strings: _StringStore, traits: _TraitPlan) -> None: strings.intern(value) +def _schema_strings( + plan: SchemaPlan, shapes: Collection[_model.Shape] +) -> frozenset[str]: + values: set[str] = set() + for shape in shapes: + values.update((shape.id.namespace, shape.id.name)) + _collect_trait_strings(values, plan.traits[shape.id]) + for member in shape.members.values(): + values.add(member.name) + _collect_trait_strings(values, plan.traits[member.id]) + return frozenset(values) + + +def _collect_trait_strings(values: set[str], traits: _TraitPlan) -> None: + if isinstance(traits, int): + return + for trait_id, value in traits.items(): + values.add(str(trait_id)) + if isinstance(value, str): + values.add(value) + + class _StringStore: __slots__ = ("_allocated", "_next", "values") diff --git a/packages/smithy-python/tests/unit/test_generation.py b/packages/smithy-python/tests/unit/test_generation.py index 7be37ca29..6065e8b4f 100644 --- a/packages/smithy-python/tests/unit/test_generation.py +++ b/packages/smithy-python/tests/unit/test_generation.py @@ -33,6 +33,15 @@ def test_types_settings_default_to_identity_selection() -> None: ) assert settings.selector == Selector.identity() assert not settings.generate_inputs_and_outputs + assert settings.schema_partition_size == 250 + assert settings.stub_partition_size == 250 + + +def test_types_settings_require_positive_partition_size() -> None: + with pytest.raises(ValueError, match="Partition sizes must be positive"): + TypesSettings("example", "1.0.0", None, schema_partition_size=0) + with pytest.raises(ValueError, match="Partition sizes must be positive"): + TypesSettings("example", "1.0.0", None, stub_partition_size=0) def test_generation_context_shares_manifest_dependencies() -> None: diff --git a/packages/smithy-python/tests/unit/test_models.py b/packages/smithy-python/tests/unit/test_models.py index 18d545c84..9b57bd399 100644 --- a/packages/smithy-python/tests/unit/test_models.py +++ b/packages/smithy-python/tests/unit/test_models.py @@ -44,6 +44,7 @@ def _context( *, service_id: str | None = "example#Service", package: str = _PACKAGE, + stub_partition_size: int | None = 250, ) -> GenerationContext[TypesSettings]: model = load_model(model_data) selection = select_type_shapes( @@ -51,7 +52,12 @@ def _context( service_id=service_id, generate_inputs_and_outputs=True, ) - settings = TypesSettings(package, "1.0.0", None) + settings = TypesSettings( + package, + "1.0.0", + None, + stub_partition_size=stub_partition_size, + ) symbols = SymbolProvider( model, selection, @@ -191,8 +197,15 @@ def _model() -> dict[str, object]: } -def _generate(tmp_path: Path) -> tuple[GenerationContext[TypesSettings], Path]: - context = _context(_model()) +def _generate( + tmp_path: Path, + *, + stub_partition_size: int | None = 250, +) -> tuple[GenerationContext[TypesSettings], Path]: + context = _context( + _model(), + stub_partition_size=stub_partition_size, + ) SchemaGenerator().generate(context) ModelGenerator().generate(context) context.files.text_file(f"src/{_PACKAGE}/__init__.py", "") @@ -433,6 +446,68 @@ def read_struct(self, schema: Any, consumer: Any) -> None: models.Original +def test_partitioned_stub_preserves_public_models( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _, source_root = _generate( + tmp_path, + stub_partition_size=2, + ) + package = source_root / _PACKAGE + facade = (package / "models.pyi").read_text() + partitions = sorted((package / "_private").glob("models_*.pyi")) + assert len(partitions) == 4 + assert "class Input(" not in facade + assert "from ._private.models_0 import (" in facade + assert "Input," in facade + assert "Input as Input" not in facade + assert "_SelectionDeserializer" not in facade + assert any("class Input(" in path.read_text() for path in partitions) + + monkeypatch.syspath_prepend( # pyright: ignore[reportUnknownMemberType] + str(source_root) + ) + models = importlib.import_module(f"{_PACKAGE}.models") + assert models.Renamed(value="value").value == "value" + assert models.SelectionText(value="value").value == "value" + + generated = tuple( + str(path) + for path in sorted(package.rglob("*")) + if path.suffix in {".py", ".pyi"} + ) + result = _run_tool("ruff", "check", "--isolated", *generated) + assert result.returncode == 0, result.stdout + result.stderr + + +@pytest.mark.parametrize(("declaration_count", "partition_count"), [(250, 0), (251, 2)]) +def test_stub_partition_threshold( + tmp_path: Path, + declaration_count: int, + partition_count: int, +) -> None: + package = "generated_threshold" + context = _context( + { + "smithy": "2.0", + "shapes": { + f"example#Shape{index}": {"type": "structure"} + for index in range(declaration_count) + }, + }, + service_id=None, + package=package, + ) + ModelGenerator().generate(context) + context.files.commit(tmp_path) + + source = tmp_path / f"src/{package}" + assert len(tuple((source / "_private").glob("models_*.pyi"))) == partition_count + facade = (source / "models.pyi").read_text() + assert ("._private.models_0" in facade) is bool(partition_count) + + def test_generated_model_names_do_not_shadow_runtime_helpers( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, @@ -797,10 +872,12 @@ def test_modeled_unit_does_not_shadow_prelude_unit( assert isinstance(decoded.action.value, models.Unit) +@pytest.mark.parametrize("stub_partition_size", [None, 2]) def test_generated_stub_passes_strict_pyright( tmp_path: Path, + stub_partition_size: int | None, ) -> None: - _, source_root = _generate(tmp_path) + _, source_root = _generate(tmp_path, stub_partition_size=stub_partition_size) (tmp_path / "check.py").write_text( ( f"from {_PACKAGE}.models import Input, Problem, Renamed, SelectionText, Tone\n\n" diff --git a/packages/smithy-python/tests/unit/test_schemas.py b/packages/smithy-python/tests/unit/test_schemas.py index ca5bc72cb..3a8b81f06 100644 --- a/packages/smithy-python/tests/unit/test_schemas.py +++ b/packages/smithy-python/tests/unit/test_schemas.py @@ -4,7 +4,9 @@ from __future__ import annotations +import importlib import runpy +import sys from collections.abc import Mapping from pathlib import Path from types import MappingProxyType @@ -32,10 +34,19 @@ from smithy_python.symbols import SymbolProvider -def _context(model_data: Mapping[str, object]) -> GenerationContext[TypesSettings]: +def _context( + model_data: Mapping[str, object], + *, + schema_partition_size: int | None = 250, +) -> GenerationContext[TypesSettings]: model = load_model(model_data) selection = select_type_shapes(model, generate_inputs_and_outputs=True) - settings = TypesSettings("example.types", "1.0.0", None) + settings = TypesSettings( + "example.types", + "1.0.0", + None, + schema_partition_size=schema_partition_size, + ) symbols = SymbolProvider(model, selection, package=settings.package) dependencies = DependencySet() return GenerationContext( @@ -231,6 +242,158 @@ def test_schema_generator_emits_lazy_hydratable_schemas(tmp_path: Path) -> None: assert str(node.members_by_index[2].expect_member_target().id) == "example#Choice" +def test_schema_generator_partitions_behind_lazy_facade( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + context = _context( + { + "smithy": "2.0", + "shapes": { + "example#Leaf": {"type": "string"}, + "example#Node": { + "type": "structure", + "members": {"leaf": {"target": "example#Leaf"}}, + }, + "example#CycleA": { + "type": "structure", + "members": {"b": {"target": "example#CycleB"}}, + }, + "example#CycleB": { + "type": "structure", + "members": {"a": {"target": "example#CycleA"}}, + }, + }, + }, + schema_partition_size=1, + ) + SchemaGenerator().generate(context) + context.files.text_file("src/example/types/__init__.py", "") + context.files.text_file("src/example/types/_private/__init__.py", "") + context.files.commit(tmp_path) + + package = tmp_path / "src/example/types/_private" + assert (package / "schemas.py").is_file() + assert (package / "schemas_0.py").is_file() + assert (package / "schemas_1.py").is_file() + facade_source = (package / "schemas.py").read_text() + assert "_PARTITION_NAMES: tuple[tuple[str, ...], ...]" in facade_source + assert "name: index" in facade_source + assert facade_source.count(".schemas_") == 1 + partitions = tuple(package.glob("schemas_*.py")) + cycle = next( + path.read_text() for path in partitions if "CYCLE_A:" in path.read_text() + ) + assert "CYCLE_B:" in cycle + + monkeypatch.syspath_prepend( # pyright: ignore[reportUnknownMemberType] + str(tmp_path / "src") + ) + facade = importlib.import_module("example.types._private.schemas") + assert "example.types._private.schemas_0" not in sys.modules + assert "example.types._private.schemas_1" not in sys.modules + + static_node = cast(tuple[Any, ...], facade.NODE) + assert facade.NODE is static_node + assert "NODE" in dir(facade) + assert "example.types._private.schemas_1" in sys.modules + assert "example.types._private.schemas_0" not in sys.modules + with pytest.raises(AttributeError): + facade.MISSING + node = _schema_compact.hydrate(static_node) + assert str(node.members["leaf"].expect_member_target().id) == "example#Leaf" + assert "example.types._private.schemas_0" in sys.modules + cycle_a = _schema_compact.hydrate(cast(tuple[Any, ...], facade.CYCLE_A)) + cycle_b = cycle_a.members["b"].expect_member_target() + assert cycle_b.members["a"].expect_member_target() is cycle_a + + +@pytest.mark.parametrize(("shape_count", "partition_count"), [(250, 0), (251, 2)]) +def test_schema_partition_threshold( + tmp_path: Path, + shape_count: int, + partition_count: int, +) -> None: + context = _context( + { + "smithy": "2.0", + "shapes": { + f"example#Shape{index}": {"type": "structure"} + for index in range(shape_count) + }, + }, + ) + SchemaGenerator().generate(context) + context.files.commit(tmp_path) + + private = tmp_path / "src/example/types/_private" + assert len(tuple(private.glob("schemas_*.py"))) == partition_count + facade = (private / "schemas.py").read_text() + assert ("_PARTITIONS" in facade) is bool(partition_count) + + +def test_schema_partitions_balance_emitted_size(tmp_path: Path) -> None: + context = _context( + { + "smithy": "2.0", + "shapes": { + "example#AHeavy": { + "type": "structure", + "members": { + f"member{index}": {"target": "smithy.api#String"} + for index in range(20) + }, + }, + "example#BTiny": {"type": "structure"}, + "example#CTiny": {"type": "structure"}, + "example#DTiny": {"type": "structure"}, + }, + }, + schema_partition_size=2, + ) + SchemaGenerator().generate(context) + context.files.commit(tmp_path) + private = tmp_path / "src/example/types/_private" + + first = (private / "schemas_0.py").read_text() + second = (private / "schemas_1.py").read_text() + assert "A_HEAVY:" in first + assert "B_TINY:" not in first + assert all(f"{name}:" in second for name in ("B_TINY", "C_TINY", "D_TINY")) + + +def test_schema_partition_weights_include_interned_strings(tmp_path: Path) -> None: + context = _context( + { + "smithy": "2.0", + "shapes": { + "example#AHeavy": { + "type": "structure", + "members": { + "value": { + "target": "smithy.api#String", + "traits": {"smithy.api#default": "x" * 10_000}, + } + }, + }, + "example#BTiny": {"type": "structure"}, + "example#CTiny": {"type": "structure"}, + "example#DTiny": {"type": "structure"}, + }, + }, + schema_partition_size=2, + ) + SchemaGenerator().generate(context) + context.files.commit(tmp_path) + private = tmp_path / "src/example/types/_private" + + first = (private / "schemas_0.py").read_text() + second = (private / "schemas_1.py").read_text() + assert "A_HEAVY:" in first + assert "B_TINY:" not in first + assert all(f"{name}:" in second for name in ("B_TINY", "C_TINY", "D_TINY")) + + def test_schema_generator_preserves_primitive_prelude_targets(tmp_path: Path) -> None: context = _context( { From b424235af06a97349492f05a71302e5e8ddc58d0 Mon Sep 17 00:00:00 2001 From: Sean O'Brien <60306702+stobrien89@users.noreply.github.com> Date: Wed, 11 Nov 2026 11:33:44 -0500 Subject: [PATCH 2/3] Match native schema partitioning to Java --- designs/codegen/partitioning.md | 20 +- designs/codegen/types.md | 2 +- ...ture-0fda37b6ad4e4dd893edbd05bbac2829.json | 4 + .../src/smithy_python/generation.py | 2 +- .../src/smithy_python/schemas.py | 334 +++++++++++++++++- .../tests/unit/test_generation.py | 2 +- .../smithy-python/tests/unit/test_schemas.py | 91 ++++- 7 files changed, 430 insertions(+), 25 deletions(-) create mode 100644 packages/smithy-python/.changes/next-release/smithy-python-feature-0fda37b6ad4e4dd893edbd05bbac2829.json diff --git a/designs/codegen/partitioning.md b/designs/codegen/partitioning.md index e4f0c8266..634394eea 100644 --- a/designs/codegen/partitioning.md +++ b/designs/codegen/partitioning.md @@ -30,7 +30,7 @@ the smaller partition files. class TypesSettings(ArtifactSettings): selector: Selector = field(default_factory=Selector.identity) generate_inputs_and_outputs: bool = False - schema_partition_size: int | None = 250 + schema_partition_size: int | None = 100 stub_partition_size: int | None = 250 ``` @@ -38,7 +38,7 @@ The CLI uses these defaults. Python callers may set either target to a positive integer or `None`. A value of `None` emits that surface as one file. Partitioning starts when the number of planned items exceeds its target. A -package with 250 schema shapes and 250 model declarations keeps the single-file +package with 100 schema shapes and 250 model declarations keeps the single-file layout. ## Output Layout @@ -93,21 +93,27 @@ Files under `._private` remain private implementation details. ## Schema Partitions -`SchemaPlan` provides schema shapes in deterministic emission order and groups -each recursive strongly connected component. A recursive group stays in one -partition so its schemas can use direct local references. +`SchemaPlan` provides schema shapes in deterministic emission order. Recursive +strongly connected components stay together so their schemas can use direct +local references. An operation's selected input and output stay together; when +operation schemas are part of the plan, the operation stays with them too. The schema partitioner: -1. Packs recursive groups in plan order using the configured shape target. +1. Packs the indivisible groups in plan order using the configured shape target. 2. Uses the resulting partition count for the final layout. 3. Estimates the emitted source size of each group. 4. Selects contiguous boundaries that balance estimated source size. +5. Uses operation reachability to move shared groups together when that reduces + the number of partitions touched by related operations. The estimate includes shape and member identifiers, schema references, and trait values. It uses plan metadata instead of rendering schemas a second time. The target determines the partition count. Final partitions may exceed the -target to balance source size, and recursive groups remain indivisible. +target to balance source size, and affinity groups remain indivisible. Affinity +moves are bounded to keep source size and shape counts within 25% of the initial +balanced layout. Shapes reached by nearly every operation are ignored so they +do not pull unrelated operations into the same partition. The facade maps each schema constant to a partition index. Its `__getattr__` imports the selected partition, reads the constant, and caches the result in diff --git a/designs/codegen/types.md b/designs/codegen/types.md index 964c4b6f8..1efa5a176 100644 --- a/designs/codegen/types.md +++ b/designs/codegen/types.md @@ -100,7 +100,7 @@ class ArtifactSettings: class TypesSettings(ArtifactSettings): selector: Selector = field(default_factory=Selector.identity) generate_inputs_and_outputs: bool = False - schema_partition_size: int | None = 250 + schema_partition_size: int | None = 100 stub_partition_size: int | None = 250 ``` diff --git a/packages/smithy-python/.changes/next-release/smithy-python-feature-0fda37b6ad4e4dd893edbd05bbac2829.json b/packages/smithy-python/.changes/next-release/smithy-python-feature-0fda37b6ad4e4dd893edbd05bbac2829.json new file mode 100644 index 000000000..8b0f0e853 --- /dev/null +++ b/packages/smithy-python/.changes/next-release/smithy-python-feature-0fda37b6ad4e4dd893edbd05bbac2829.json @@ -0,0 +1,4 @@ +{ + "type": "feature", + "description": "Partition large generated schema and model stub files to reduce memory usage and static analysis time." +} \ No newline at end of file diff --git a/packages/smithy-python/src/smithy_python/generation.py b/packages/smithy-python/src/smithy_python/generation.py index b82369040..de29ad6d3 100644 --- a/packages/smithy-python/src/smithy_python/generation.py +++ b/packages/smithy-python/src/smithy_python/generation.py @@ -44,7 +44,7 @@ class TypesSettings(ArtifactSettings): selector: _selection.Selector = field(default_factory=_selection.Selector.identity) generate_inputs_and_outputs: bool = False - schema_partition_size: int | None = 250 + schema_partition_size: int | None = 100 stub_partition_size: int | None = 250 def __post_init__(self) -> None: diff --git a/packages/smithy-python/src/smithy_python/schemas.py b/packages/smithy-python/src/smithy_python/schemas.py index cf0764641..e386b2d1f 100644 --- a/packages/smithy-python/src/smithy_python/schemas.py +++ b/packages/smithy-python/src/smithy_python/schemas.py @@ -6,7 +6,7 @@ from collections.abc import Collection, Mapping from dataclasses import dataclass -from math import isfinite +from math import ceil, isfinite, sqrt from types import MappingProxyType from typing import TYPE_CHECKING, Final @@ -21,6 +21,10 @@ __all__ = ["SchemaGenerator", "SchemaPlan"] +_PARTITION_BALANCE_TOLERANCE: Final = 0.25 +_MIN_COMMON_OPERATION_COUNT: Final = 12 +_AFFINITY_PASSES: Final = 3 + @dataclass(frozen=True, slots=True) class SchemaPlan: @@ -28,6 +32,7 @@ class SchemaPlan: shapes: tuple[_model.Shape, ...] partition_groups: tuple[tuple[_model.Shape, ...], ...] + operations_by_group: tuple[frozenset[_model.ShapeId], ...] names: Mapping[_model.ShapeId, str] strings: Mapping[str, str] traits: Mapping[_model.ShapeId, _TraitPlan] @@ -40,7 +45,13 @@ def from_shapes( symbols: _symbols.SymbolProvider, ) -> SchemaPlan: """Plan schema order, filtered traits, names, and interned strings.""" - emitted, partition_groups = _schema_order(shapes) + emitted, recursive_groups = _schema_order(shapes) + operations = _partition_operations(model, emitted) + partition_groups = _merge_operation_groups(operations, recursive_groups) + operations_by_group = _operations_by_group( + partition_groups, + _operation_closures(emitted, operations), + ) included = _included_traits(model) inputs = { operation.input for operation in model.iter_shapes(_model.OperationShape) @@ -71,6 +82,7 @@ def from_shapes( return cls( shapes=emitted, partition_groups=partition_groups, + operations_by_group=operations_by_group, names=MappingProxyType( {shape.id: symbols.schema_name(shape.id) for shape in emitted} ), @@ -173,9 +185,9 @@ def _write_schema_module( ) -> None: local = frozenset(shape.id for shape in shapes) if facade is not None and any( - member.target in plan.names and member.target not in local + target in plan.names and target not in local for shape in shapes - for member in shape.members.values() + for target in _schema_references(shape) ): writer.line("from . import ", facade) writer.line() @@ -542,7 +554,7 @@ def _schema_order( positions = {shape.id: index for index, shape in enumerate(shapes)} graph = { shape.id: tuple( - member.target for member in shape.members.values() if member.target in by_id + target for target in _schema_references(shape) if target in by_id ) for shape in shapes } @@ -610,6 +622,61 @@ def _schema_order( return emitted, tuple(groups) +def _schema_references(shape: _model.Shape) -> tuple[_model.ShapeId, ...]: + references = dict.fromkeys(member.target for member in shape.members.values()) + if isinstance(shape, _model.OperationShape): + references.update(dict.fromkeys((shape.input, shape.output, *shape.errors))) + return tuple(references) + + +def _merge_operation_groups( + operations: tuple[_model.OperationShape, ...], + groups: tuple[tuple[_model.Shape, ...], ...], +) -> tuple[tuple[_model.Shape, ...], ...]: + group_by_shape: dict[_model.ShapeId, int] = {} + parents = list(range(len(groups))) + for index, group in enumerate(groups): + for shape in group: + group_by_shape[shape.id] = index + + for operation in operations: + _merge_groups(parents, group_by_shape, operation.id, operation.input) + _merge_groups(parents, group_by_shape, operation.id, operation.output) + _merge_groups(parents, group_by_shape, operation.input, operation.output) + + merged: dict[int, list[_model.Shape]] = {} + for index, group in enumerate(groups): + merged.setdefault(_find_group(parents, index), []).extend(group) + return tuple(tuple(group) for group in merged.values()) + + +def _merge_groups( + parents: list[int], + group_by_shape: Mapping[_model.ShapeId, int], + left_shape: _model.ShapeId, + right_shape: _model.ShapeId, +) -> None: + left = group_by_shape.get(left_shape) + right = group_by_shape.get(right_shape) + if left is None or right is None: + return + left_root = _find_group(parents, left) + right_root = _find_group(parents, right) + if left_root != right_root: + parents[max(left_root, right_root)] = min(left_root, right_root) + + +def _find_group(parents: list[int], group: int) -> int: + root = group + while parents[root] != root: + root = parents[root] + while parents[group] != group: + next_group = parents[group] + parents[group] = root + group = next_group + return root + + def _schema_partitions( plan: SchemaPlan, size: int, @@ -629,27 +696,262 @@ def _schema_partitions( return tuple(partitions) weights = tuple(_schema_group_weight(plan, group) for group in groups) - partitions.clear() + layout = _source_balanced_layout(weights, partition_count) + layout = _improve_operation_affinity_layout( + tuple(map(len, groups)), + weights, + plan.operations_by_group, + layout, + size, + ) + return tuple( + tuple(shape for group_index in partition for shape in groups[group_index]) + for partition in layout + ) + + +def _source_balanced_layout( + weights: tuple[int, ...], + partition_count: int, +) -> list[list[int]]: + partitions: list[list[int]] = [] start = 0 remaining_weight = sum(weights) for remaining_parts in range(partition_count, 1, -1): target = remaining_weight / remaining_parts end = start weight = 0 - last_end = len(groups) - remaining_parts + 1 + last_end = len(weights) - remaining_parts + 1 while end < last_end: next_weight = weight + weights[end] if weight and abs(target - weight) <= abs(target - next_weight): break weight = next_weight end += 1 - partitions.append( - tuple(shape for group in groups[start:end] for shape in group) - ) + partitions.append(list(range(start, end))) start = end remaining_weight -= weight - partitions.append(tuple(shape for group in groups[start:] for shape in group)) - return tuple(partitions) + partitions.append(list(range(start, len(weights)))) + return partitions + + +def _operation_closures( + shapes: tuple[_model.Shape, ...], + operations: tuple[_model.OperationShape, ...], +) -> dict[_model.ShapeId, frozenset[_model.ShapeId]]: + by_id = {shape.id: shape for shape in shapes} + closures: dict[_model.ShapeId, frozenset[_model.ShapeId]] = {} + for operation in operations: + visited: set[_model.ShapeId] = set() + pending = [operation.id] + while pending: + current = pending.pop() + if current in visited: + continue + visited.add(current) + current_shape = operation if current == operation.id else by_id.get(current) + if current_shape is None: + continue + for target in reversed(_schema_references(current_shape)): + if target in by_id and target not in visited: + pending.append(target) + closures[operation.id] = frozenset(visited) + return closures + + +def _partition_operations( + model: _model.Model, + shapes: tuple[_model.Shape, ...], +) -> tuple[_model.OperationShape, ...]: + emitted = tuple( + shape for shape in shapes if isinstance(shape, _model.OperationShape) + ) + if emitted: + return emitted + return tuple(model.iter_shapes(_model.OperationShape)) + + +def _operations_by_group( + groups: tuple[tuple[_model.Shape, ...], ...], + operation_closures: Mapping[_model.ShapeId, Collection[_model.ShapeId]], +) -> tuple[frozenset[_model.ShapeId], ...]: + group_by_shape = { + shape.id: index for index, group in enumerate(groups) for shape in group + } + operations_by_group: list[set[_model.ShapeId]] = [set() for _ in groups] + for operation, closure in operation_closures.items(): + for shape_id in closure: + if (group := group_by_shape.get(shape_id)) is not None: + operations_by_group[group].add(operation) + return tuple(frozenset(operations) for operations in operations_by_group) + + +def _improve_operation_affinity_layout( + group_shape_counts: tuple[int, ...], + group_weights: tuple[int, ...], + group_operations: tuple[Collection[_model.ShapeId], ...], + initial_layout: list[list[int]], + target_size: int, +) -> list[list[int]]: + assigned = [list(partition) for partition in initial_layout] + partition_count = len(assigned) + shape_counts = [0] * partition_count + source_weights = [0] * partition_count + operation_counts: list[dict[_model.ShapeId, int]] = [] + operation_locations: dict[_model.ShapeId, set[int]] = {} + for partition, group_indexes in enumerate(assigned): + counts: dict[_model.ShapeId, int] = {} + for group_index in group_indexes: + shape_counts[partition] += group_shape_counts[group_index] + source_weights[partition] += group_weights[group_index] + for operation in sorted(group_operations[group_index]): + counts[operation] = counts.get(operation, 0) + 1 + operation_counts.append(counts) + for operation in counts: + operation_locations.setdefault(operation, set()).add(partition) + + ideal_source_weight = sum(group_weights) / partition_count + ideal_shape_count = sum(group_shape_counts) / partition_count + source_ceilings = [ + max( + ceil(ideal_source_weight * (1 + _PARTITION_BALANCE_TOLERANCE)), + source_weight, + ) + for source_weight in source_weights + ] + source_floors = [ + min( + int(ideal_source_weight * (1 - _PARTITION_BALANCE_TOLERANCE)), + source_weight, + ) + for source_weight in source_weights + ] + shape_ceilings = [ + max( + shape_count, + ceil( + max(target_size, ideal_shape_count) * (1 + _PARTITION_BALANCE_TOLERANCE) + ), + ) + for shape_count in shape_counts + ] + common_threshold = max( + _MIN_COMMON_OPERATION_COUNT, + ceil(sqrt(len(operation_locations))), + ) + + for _ in range(_AFFINITY_PASSES): + changed = False + for source in range(partition_count): + for group_index in tuple(assigned[source]): + if len(assigned[source]) == 1: + continue + operations = group_operations[group_index] + if not 2 <= len(operations) <= common_threshold: + continue + candidates = sorted( + { + partition + for operation in sorted(operations) + for partition in operation_locations[operation] + if partition != source + } + ) + best: tuple[int, float, int] | None = None + for target in candidates: + new_source_weight = ( + source_weights[source] - group_weights[group_index] + ) + new_target_weight = ( + source_weights[target] + group_weights[group_index] + ) + new_target_shapes = ( + shape_counts[target] + group_shape_counts[group_index] + ) + if ( + new_source_weight < source_floors[source] + or new_target_weight > source_ceilings[target] + or new_target_shapes > shape_ceilings[target] + ): + continue + delta = _operation_touch_delta( + operations, + operation_counts[source], + operation_counts[target], + ) + if delta >= 0: + continue + balance = abs(new_source_weight - ideal_source_weight) + abs( + new_target_weight - ideal_source_weight + ) + candidate = (delta, balance, target) + if best is None or candidate < best: + best = candidate + if best is None: + continue + + target = best[2] + assigned[source].remove(group_index) + assigned[target].append(group_index) + shape_counts[source] -= group_shape_counts[group_index] + shape_counts[target] += group_shape_counts[group_index] + source_weights[source] -= group_weights[group_index] + source_weights[target] += group_weights[group_index] + _update_operation_counts( + source, + operations, + operation_counts, + operation_locations, + -1, + ) + _update_operation_counts( + target, + operations, + operation_counts, + operation_locations, + 1, + ) + changed = True + if not changed: + break + + return [sorted(partition) for partition in assigned] + + +def _operation_touch_delta( + operations: Collection[_model.ShapeId], + source_counts: Mapping[_model.ShapeId, int], + target_counts: Mapping[_model.ShapeId, int], +) -> int: + delta = 0 + for operation in sorted(operations): + if target_counts.get(operation, 0) == 0: + delta += 1 + if source_counts.get(operation, 0) == 1: + delta -= 1 + return delta + + +def _update_operation_counts( + partition: int, + operations: Collection[_model.ShapeId], + operation_counts: list[dict[_model.ShapeId, int]], + operation_locations: Mapping[_model.ShapeId, set[int]], + adjustment: int, +) -> None: + counts = operation_counts[partition] + for operation in sorted(operations): + current = counts.get(operation, 0) + updated = current + adjustment + if updated == 0: + counts.pop(operation) + else: + counts[operation] = updated + locations = operation_locations[operation] + if current == 0 and updated == 1: + locations.add(partition) + elif current == 1 and updated == 0: + locations.remove(partition) def _schema_group_weight(plan: SchemaPlan, group: tuple[_model.Shape, ...]) -> int: @@ -662,7 +964,7 @@ def _schema_group_weight(plan: SchemaPlan, group: tuple[_model.Shape, ...]) -> i def _schema_shape_weight(plan: SchemaPlan, shape: _model.Shape) -> int: """Estimate compact-schema source bytes without rendering shapes twice.""" - return ( + weight = ( 60 + len(shape.id.namespace) + 2 * len(shape.id.name) @@ -675,6 +977,12 @@ def _schema_shape_weight(plan: SchemaPlan, shape: _model.Shape) -> int: for member in shape.members.values() ) ) + if isinstance(shape, _model.OperationShape): + weight += sum( + len(plan.names.get(target, target.name)) + for target in (shape.input, shape.output, *shape.errors) + ) + return weight def _source_size(parts: Collection[_SourcePart]) -> int: diff --git a/packages/smithy-python/tests/unit/test_generation.py b/packages/smithy-python/tests/unit/test_generation.py index 6065e8b4f..9a5ef2f91 100644 --- a/packages/smithy-python/tests/unit/test_generation.py +++ b/packages/smithy-python/tests/unit/test_generation.py @@ -33,7 +33,7 @@ def test_types_settings_default_to_identity_selection() -> None: ) assert settings.selector == Selector.identity() assert not settings.generate_inputs_and_outputs - assert settings.schema_partition_size == 250 + assert settings.schema_partition_size == 100 assert settings.stub_partition_size == 250 diff --git a/packages/smithy-python/tests/unit/test_schemas.py b/packages/smithy-python/tests/unit/test_schemas.py index 3a8b81f06..154617162 100644 --- a/packages/smithy-python/tests/unit/test_schemas.py +++ b/packages/smithy-python/tests/unit/test_schemas.py @@ -28,6 +28,7 @@ _TYPE_INTS, _UNIT_INT, SchemaGenerator, + _improve_operation_affinity_layout, _StringStore, ) from smithy_python.selection import select_type_shapes @@ -37,7 +38,7 @@ def _context( model_data: Mapping[str, object], *, - schema_partition_size: int | None = 250, + schema_partition_size: int | None = 100, ) -> GenerationContext[TypesSettings]: model = load_model(model_data) selection = select_type_shapes(model, generate_inputs_and_outputs=True) @@ -308,7 +309,7 @@ def test_schema_generator_partitions_behind_lazy_facade( assert cycle_b.members["a"].expect_member_target() is cycle_a -@pytest.mark.parametrize(("shape_count", "partition_count"), [(250, 0), (251, 2)]) +@pytest.mark.parametrize(("shape_count", "partition_count"), [(100, 0), (101, 2)]) def test_schema_partition_threshold( tmp_path: Path, shape_count: int, @@ -332,6 +333,92 @@ def test_schema_partition_threshold( assert ("_PARTITIONS" in facade) is bool(partition_count) +def test_schema_partition_groups_keep_operation_input_and_output_together() -> None: + context = _context( + { + "smithy": "2.0", + "shapes": { + "example#Input": {"type": "structure"}, + "example#Output": {"type": "structure"}, + "example#Operation": { + "type": "operation", + "input": {"target": "example#Input"}, + "output": {"target": "example#Output"}, + }, + }, + } + ) + + groups = context.plan.schemas.partition_groups + operations_by_group = context.plan.schemas.operations_by_group + + group_index = next( + index + for index, group in enumerate(groups) + if {str(shape.id) for shape in group} == {"example#Input", "example#Output"} + ) + assert tuple(map(str, operations_by_group[group_index])) == ("example#Operation",) + + +def test_operation_affinity_moves_shared_shapes_toward_related_operations() -> None: + create = ShapeId.from_string("example.affinity#Create") + describe = ShapeId.from_string("example.affinity#Describe") + unrelated = ShapeId.from_string("example.affinity#Unrelated") + + optimized = _improve_operation_affinity_layout( + (1, 1, 1, 1), + (5, 5, 10, 5), + ( + frozenset((create,)), + frozenset((create, describe)), + frozenset((unrelated,)), + frozenset((describe,)), + ), + [[0], [1, 2], [3]], + 2, + ) + + assert optimized == [[0, 1], [2], [3]] + + +def test_operation_affinity_is_independent_of_operation_iteration_order() -> None: + create = ShapeId.from_string("example.affinity#Create") + describe = ShapeId.from_string("example.affinity#Describe") + update = ShapeId.from_string("example.affinity#Update") + shape_counts = (1, 1, 1, 1, 1) + weights = (5, 5, 5, 5, 5) + initial_layout = [[0, 1], [2, 3], [4]] + + forward = _improve_operation_affinity_layout( + shape_counts, + weights, + ( + (create, describe), + (update,), + (describe, create), + (describe, update), + (update, create), + ), + initial_layout, + 2, + ) + reversed_operations = _improve_operation_affinity_layout( + shape_counts, + weights, + ( + (describe, create), + (update,), + (create, describe), + (update, describe), + (create, update), + ), + initial_layout, + 2, + ) + + assert forward == reversed_operations + + def test_schema_partitions_balance_emitted_size(tmp_path: Path) -> None: context = _context( { From 2240be8c5f39027a9736f260a04b6ce3f537bf1b Mon Sep 17 00:00:00 2001 From: Sean O'Brien <60306702+stobrien89@users.noreply.github.com> Date: Wed, 11 Nov 2026 14:45:06 -0500 Subject: [PATCH 3/3] Fix operation closure lint --- packages/smithy-python/src/smithy_python/schemas.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/packages/smithy-python/src/smithy_python/schemas.py b/packages/smithy-python/src/smithy_python/schemas.py index e386b2d1f..af565ea38 100644 --- a/packages/smithy-python/src/smithy_python/schemas.py +++ b/packages/smithy-python/src/smithy_python/schemas.py @@ -752,9 +752,11 @@ def _operation_closures( current_shape = operation if current == operation.id else by_id.get(current) if current_shape is None: continue - for target in reversed(_schema_references(current_shape)): - if target in by_id and target not in visited: - pending.append(target) + pending.extend( + target + for target in reversed(_schema_references(current_shape)) + if target in by_id and target not in visited + ) closures[operation.id] = frozenset(visited) return closures