Generate external transform wrappers using a script

[ahmedabu98]

Implementing a script that generates wrappers for external SchemaTransforms, according to Option #3 in the following design doc: https://s.apache.org/autogen-wrappers

The script's workflow takes place in setup.py, which can be invoked for local setup or for building the SDK for a Beam release. Files are generated in a subdirectory that is ignored by git. From there, the wrappers can be imported into relevant __init__.py files.

Wrappers are generated along with any documentation provided from the underlying SchemaTransform, and is in compliance with existing linting and pydoc rules.

With the expansion service YAML config, the following transform YAML config is generated:

Transform YAML config

- default_service: sdks:java:io:expansion-service:shadowJar
  description: 'Outputs a PCollection of Beam Rows, each containing a single INT64
    number called "value". The count is produced from the given "start"value and either
    up to the given "end" or until 2^63 - 1.

    To produce an unbounded PCollection, simply do not specify an "end" value. Unbounded
    sequences can specify a "rate" for output elements.

    In all cases, the sequence of numbers is generated in parallel, so there is no
    inherent ordering between the generated values'
  destinations:
    python: apache_beam/io
  fields:
    end:
      description: The maximum number to generate (exclusive). Will be an unbounded
        sequence if left unspecified.
      nullable: true
      type: numpy.int64
    rate:
      description: Specifies the rate to generate a given number of elements per a
        given number of seconds. Applicable only to unbounded sequences.
      nullable: true
      type: Row(seconds=typing.Union[numpy.int64, NoneType], elements=<class 'numpy.int64'>)
    start:
      description: The minimum number to generate (inclusive).
      nullable: false
      type: numpy.int64
  identifier: beam:schematransform:org.apache.beam:generate_sequence:v1
  name: GenerateSequence

From this config, external transform wrappers are created and stored in appropriate modules. For example, our config gives us the following apache_beam/transforms/xlang/io.py file:

GenerateSequence wrapper

# NOTE: This file contains autogenerated external transform(s)
# and should not be edited by hand.
# Refer to gen_xlang_wrappers.py for more info.

"""Cross-language transforms in this module can be imported from the
:py:mod:`apache_beam.io` package."""

# pylint:disable=line-too-long

from apache_beam.transforms.external import BeamJarExpansionService
from apache_beam.transforms.external_transform_provider import ExternalTransform


class GenerateSequence(ExternalTransform):
  """
  Outputs a PCollection of Beam Rows, each containing a single INT64 number
  called "value". The count is produced from the given "start" value and either
  up to the given "end" or until 2^63 - 1.
  To produce an unbounded PCollection, simply do not specify an "end" value.
  Unbounded sequences can specify a "rate" for output elements.
  In all cases, the sequence of numbers is generated in parallel, so there is no
  inherent ordering between the generated values
  """
  identifier = "beam:schematransform:org.apache.beam:generate_sequence:v1"

  def __init__(self, start, end=None, rate=None, expansion_service=None):
    """
    :param start: (numpy.int64)
      The minimum number to generate (inclusive). 
    :param end: (numpy.int64)
      The maximum number to generate (exclusive). Will be an unbounded
      sequence if left unspecified. 
    :param rate: (Row(seconds=typing.Union[numpy.int64, NoneType], elements=<class 'numpy.int64'>))
      Specifies the rate to generate a given number of elements per a given
      number of seconds. Applicable only to unbounded sequences. 
    """
    self.default_expansion_service = BeamJarExpansionService(
        "sdks:java:io:expansion-service:shadowJar")
    super().__init__(
        start=start, end=end, rate=rate, expansion_service=expansion_service)

---

Including documentation for how the script is used, and unit tests for different parts of the script.

Adding a gradle command ./gradlew generateExternalTransformsConfig to build the configs (it takes care of building the relevant expansion jars beforehand). Once th

Also adding a PreCommit test that generates the transform config from scratch and compares it with the existing one. This serves to indicate whether a change will render the existing config out of sync. To resolve, one would re-generate the config (with ./gradlew generateExternalTransformsConfig) and commit the changes.

[github-actions]

Assigning reviewers. If you would like to opt out of this review, comment assign to next reviewer:

R: @jrmccluskey for label python.
R: @damccorm for label build.

Available commands:

• stop reviewer notifications - opt out of the automated review tooling
• remind me after tests pass - tag the comment author after tests pass
• waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).

[damccorm]

R: @robertwb @tvalentyn @chamikaramj

(manually requesting to make the review bot happy)

[github-actions]

Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control

[robertwb]

If we don't do this now, could you at least file a bug and drop a TODO.

[robertwb]

Ah, yes, that's right. Hopefully the file generation itself isn't *that* slow.

…

On Mon, Feb 12, 2024 at 3:02 PM Ahmed Abualsaud ***@***.***> wrote: ***@***.**** commented on this pull request. ------------------------------ In sdks/python/setup.py <#29834 (comment)>: > + # if exists, this directory will have at least its __init__.py file + if (not os.path.exists(generated_transforms_dir) or + len(os.listdir(generated_transforms_dir)) <= 1): + message = 'External transform wrappers have not been generated ' + if not script_exists: + message += 'and the generation script `gen_xlang_wrappers.py`' + if not config_exists: + message += 'and the standard external transforms config' + message += ' could not be found' + warnings.warn(message) + else: + warnings.warn( + 'Skipping external transform wrapper generation as they ' + 'are already generated.') + return + out = subprocess.run([ This will make pretty much every run of setup.py slow, right? Not anymore, since we decided to decouple the generation steps. setup.py expects the transform config to already exist, so no expansion + discovery is done here. All it does is generate files based on the existing config. — Reply to this email directly, view it on GitHub <#29834 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AADWVAJPRS6PTJHJGTPTHE3YTKNPZAVCNFSM6AAAAABA4ZG2QSVHI2DSMVQWIX3LMV43YUDVNRWFEZLROVSXG5CSMV3GSZLXHMYTQNZWGUYTSNZSGA> . You are receiving this because you were mentioned.Message ID: ***@***.***>

[tvalentyn]

I think we need to put the import back in the try:, except: pass block?

That will work, alternatively, you can have dedicated functions: to generate config, and to generate wrappers, and make necessary imports only in functions where imports are required; you might have to add # pylint: disable=g-import-not-at-top.

[ahmedabu98]

@tvalentyn this is already the case, the steps are in different functions. e.g. we're only importing apache_beam when using the function to generate transform configs.
But this function is unusable in a clean setup (one that doesn't have generated modules) unless we make imports like from apache_beam.transforms.xlang.io import * optional.

[ahmedabu98]

The newly added precommit that performs a transform config sync check (beam_PreCommit_Xlang_Generated_Transforms) is running green on my fork: https://github.com/ahmedabu98/beam/actions/runs/7933231703/

Will merge after current tests pass