github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/examples/complete/juliaset/setup.py (about)

     1  #
     2  # Licensed to the Apache Software Foundation (ASF) under one or more
     3  # contributor license agreements.  See the NOTICE file distributed with
     4  # this work for additional information regarding copyright ownership.
     5  # The ASF licenses this file to You under the Apache License, Version 2.0
     6  # (the "License"); you may not use this file except in compliance with
     7  # the License.  You may obtain a copy of the License at
     8  #
     9  #    http://www.apache.org/licenses/LICENSE-2.0
    10  #
    11  # Unless required by applicable law or agreed to in writing, software
    12  # distributed under the License is distributed on an "AS IS" BASIS,
    13  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14  # See the License for the specific language governing permissions and
    15  # limitations under the License.
    16  #
    17  
    18  """Setup.py module for the workflow's worker utilities.
    19  
    20  All the workflow related code is gathered in a package that will be built as a
    21  source distribution, staged in the staging area for the workflow being run and
    22  then installed in the workers when they start running.
    23  
    24  This behavior is triggered by specifying the --setup_file command line option
    25  when running the workflow for remote execution.
    26  """
    27  
    28  # pytype: skip-file
    29  
    30  import subprocess
    31  
    32  import setuptools
    33  
    34  # It is recommended to import setuptools prior to importing distutils to avoid
    35  # using legacy behavior from distutils.
    36  # https://setuptools.readthedocs.io/en/latest/history.html#v48-0-0
    37  from distutils.command.build import build as _build  # isort:skip
    38  
    39  
    40  # This class handles the pip install mechanism.
    41  class build(_build):  # pylint: disable=invalid-name
    42    """A build command class that will be invoked during package install.
    43  
    44    The package built using the current setup.py will be staged and later
    45    installed in the worker using `pip install package'. This class will be
    46    instantiated during install for this specific scenario and will trigger
    47    running the custom commands specified.
    48    """
    49    sub_commands = _build.sub_commands + [('CustomCommands', None)]
    50  
    51  
    52  # Some custom command to run during setup. The command is not essential for this
    53  # workflow. It is used here as an example. Each command will spawn a child
    54  # process. Typically, these commands will include steps to install non-Python
    55  # packages. For instance, to install a C++-based library libjpeg62 the following
    56  # two commands will have to be added:
    57  #
    58  #     ['apt-get', 'update'],
    59  #     ['apt-get', '--assume-yes', 'install', 'libjpeg62'],
    60  #
    61  # First, note that there is no need to use the sudo command because the setup
    62  # script runs with appropriate access.
    63  # Second, if apt-get tool is used then the first command needs to be 'apt-get
    64  # update' so the tool refreshes itself and initializes links to download
    65  # repositories.  Without this initial step the other apt-get install commands
    66  # will fail with package not found errors. Note also --assume-yes option which
    67  # shortcuts the interactive confirmation.
    68  #
    69  # Note that in this example custom commands will run after installing required
    70  # packages. If you have a PyPI package that depends on one of the custom
    71  # commands, move installation of the dependent package to the list of custom
    72  # commands, e.g.:
    73  #
    74  #     ['pip', 'install', 'my_package'],
    75  #
    76  # TODO(https://github.com/apache/beam/issues/18568): Output from the custom
    77  # commands are missing from the logs. The output of custom commands (including
    78  # failures) will be logged in the worker-startup log.
    79  CUSTOM_COMMANDS = [['echo', 'Custom command worked!']]
    80  
    81  
    82  class CustomCommands(setuptools.Command):
    83    """A setuptools Command class able to run arbitrary commands."""
    84    def initialize_options(self):
    85      pass
    86  
    87    def finalize_options(self):
    88      pass
    89  
    90    def RunCustomCommand(self, command_list):
    91      print('Running command: %s' % command_list)
    92      p = subprocess.Popen(
    93          command_list,
    94          stdin=subprocess.PIPE,
    95          stdout=subprocess.PIPE,
    96          stderr=subprocess.STDOUT)
    97      # Can use communicate(input='y\n'.encode()) if the command run requires
    98      # some confirmation.
    99      stdout_data, _ = p.communicate()
   100      print('Command output: %s' % stdout_data)
   101      if p.returncode != 0:
   102        raise RuntimeError(
   103            'Command %s failed: exit code: %s' % (command_list, p.returncode))
   104  
   105    def run(self):
   106      for command in CUSTOM_COMMANDS:
   107        self.RunCustomCommand(command)
   108  
   109  
   110  # Configure the required packages and scripts to install.
   111  # Note that the Python Dataflow containers come with numpy already installed
   112  # so this dependency will not trigger anything to be installed unless a version
   113  # restriction is specified.
   114  REQUIRED_PACKAGES = [
   115      'numpy',
   116  ]
   117  
   118  setuptools.setup(
   119      name='juliaset',
   120      version='0.0.1',
   121      description='Julia set workflow package.',
   122      install_requires=REQUIRED_PACKAGES,
   123      packages=setuptools.find_packages(),
   124      cmdclass={
   125          # Command class instantiated and run during pip install scenarios.
   126          'build': build,
   127          'CustomCommands': CustomCommands,
   128      })