Undo #35962 and move from distutils to setuptools instead

sdks/python/apache_beam/examples/complete/juliaset/juliaset/juliaset_test_it.py

@@ -38,7 +38,7 @@
 class JuliaSetTestIT(unittest.TestCase):
   GRID_SIZE = 1000
 
-  def test_run_example_with_requirements_file(self):
+  def test_run_example_with_setup_file(self):
     pipeline = TestPipeline(is_integration_test=True)
     coordinate_output = FileSystems.join(
         pipeline.get_option('output'),
@@ -47,8 +47,8 @@ def test_run_example_with_requirements_file(self):
     extra_args = {
         'coordinate_output': coordinate_output,
         'grid_size': self.GRID_SIZE,
-        'requirements_file': os.path.normpath(
-            os.path.join(os.path.dirname(__file__), '..', 'requirements.txt')),
+        'setup_file': os.path.normpath(
+            os.path.join(os.path.dirname(__file__), '..', 'setup.py')),
         'on_success_matcher': all_of(PipelineStateMatcher(PipelineState.DONE)),
     }
     args = pipeline.get_full_options_as_args(**extra_args)

sdks/python/apache_beam/examples/complete/juliaset/juliaset_main.py

@@ -21,12 +21,17 @@
 workflow. It is organized in this way so that it can be packaged as a Python
 package and later installed in the VM workers executing the job. The root
 directory for the example contains just a "driver" script to launch the job
-and the requirements.txt file needed to create a package.
+and the setup.py file needed to create a package.
 
 The advantages for organizing the code is that large projects will naturally
 evolve beyond just one module and you will have to make sure the additional
 modules are present in the worker.
 
+In Python Dataflow, using the --setup_file option when submitting a job, will
+trigger creating a source distribution (as if running python setup.py sdist) and
+then staging the resulting tarball in the staging area. The workers, upon
+startup, will install the tarball.
+
 Below is a complete command line for running the juliaset workflow remotely as
 an example:
 
@@ -35,7 +40,7 @@
   --project YOUR-PROJECT \
   --region GCE-REGION \
   --runner DataflowRunner \
-  --requirements_file ./requirements.txt \
+  --setup_file ./setup.py \
   --staging_location gs://YOUR-BUCKET/juliaset/staging \
   --temp_location gs://YOUR-BUCKET/juliaset/temp \
   --coordinate_output gs://YOUR-BUCKET/juliaset/out \

sdks/python/apache_beam/examples/complete/juliaset/setup.py

@@ -0,0 +1,125 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Setup.py module for the workflow's worker utilities.
+
+All the workflow related code is gathered in a package that will be built as a
+source distribution, staged in the staging area for the workflow being run and
+then installed in the workers when they start running.
+
+This behavior is triggered by specifying the --setup_file command line option
+when running the workflow for remote execution.
+"""
+
+# pytype: skip-file
+
+import subprocess
+
+import setuptools
+
+from setuptools.command.build import build as _build  # isort:skip
+
+
+# This class handles the pip install mechanism.
+class build(_build):  # pylint: disable=invalid-name
+  """A build command class that will be invoked during package install.
+
+  The package built using the current setup.py will be staged and later
+  installed in the worker using `pip install package'. This class will be
+  instantiated during install for this specific scenario and will trigger
+  running the custom commands specified.
+  """
+  sub_commands = _build.sub_commands + [('CustomCommands', None)]
+
+
+# Some custom command to run during setup. The command is not essential for this
+# workflow. It is used here as an example. Each command will spawn a child
+# process. Typically, these commands will include steps to install non-Python
+# packages. For instance, to install a C++-based library libjpeg62 the following
+# two commands will have to be added:
+#
+#     ['apt-get', 'update'],
+#     ['apt-get', '--assume-yes', 'install', 'libjpeg62'],
+#
+# First, note that there is no need to use the sudo command because the setup
+# script runs with appropriate access.
+# Second, if apt-get tool is used then the first command needs to be 'apt-get
+# update' so the tool refreshes itself and initializes links to download
+# repositories.  Without this initial step the other apt-get install commands
+# will fail with package not found errors. Note also --assume-yes option which
+# shortcuts the interactive confirmation.
+#
+# Note that in this example custom commands will run after installing required
+# packages. If you have a PyPI package that depends on one of the custom
+# commands, move installation of the dependent package to the list of custom
+# commands, e.g.:
+#
+#     ['pip', 'install', 'my_package'],
+#
+# TODO(https://github.com/apache/beam/issues/18568): Output from the custom
+# commands are missing from the logs. The output of custom commands (including
+# failures) will be logged in the worker-startup log.
+CUSTOM_COMMANDS = [['echo', 'Custom command worked!']]
+
+
+class CustomCommands(setuptools.Command):
+  """A setuptools Command class able to run arbitrary commands."""
+  def initialize_options(self):
+    pass
+
+  def finalize_options(self):
+    pass
+
+  def RunCustomCommand(self, command_list):
+    print('Running command: %s' % command_list)
+    p = subprocess.Popen(
+        command_list,
+        stdin=subprocess.PIPE,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT)
+    # Can use communicate(input='y\n'.encode()) if the command run requires
+    # some confirmation.
+    stdout_data, _ = p.communicate()
+    print('Command output: %s' % stdout_data)
+    if p.returncode != 0:
+      raise RuntimeError(
+          'Command %s failed: exit code: %s' % (command_list, p.returncode))
+
+  def run(self):
+    for command in CUSTOM_COMMANDS:
+      self.RunCustomCommand(command)
+
+
+# Configure the required packages and scripts to install.
+# Note that the Python Dataflow containers come with numpy already installed
+# so this dependency will not trigger anything to be installed unless a version
+# restriction is specified.
+REQUIRED_PACKAGES = [
+    'numpy',
+]
+
+setuptools.setup(
+    name='juliaset',
+    version='0.0.1',
+    description='Julia set workflow package.',
+    install_requires=REQUIRED_PACKAGES,
+    packages=setuptools.find_packages(),
+    cmdclass={
+        # Command class instantiated and run during pip install scenarios.
+        'build': build,
+        'CustomCommands': CustomCommands,
+    })