Add logic for building "ml-pipelines-sdk" pipeline SDK package with minimal dependencies.

PiperOrigin-RevId: 345747124

Author: tfx-copybara

README.ml-pipelines-sdk.md

@@ -0,0 +1,12 @@
+<!-- See: www.tensorflow.org/tfx/ -->
+
+# ML Pipelines SDK
+
+The `ml-pipelines-sdk` package provides the pipeline-authoring SDK for the
+[TFX](https://tensorflow.org/tfx) machine learning framework. This package is
+a subset of the `tfx` package and is provided for TFX pipeline authors who wish
+to minimize runtime dependencies.
+
+Releases are synchronized in lock-step with the
+[`tfx` package](https://pypi.org/project/tfx) and overall release notes
+can be found at https://github.com/tensorflow/tfx/releases.

package_build/README.md

@@ -0,0 +1,52 @@
+# TFX package
+
+TFX is packaged as the `tfx` package on PyPI. We recommend that users install
+TFX using `pip install tfx`. As of version 0.26.0, users also have the option
+to install a standalone version of the TFX pipeline authoring SDK, as the
+`ml-pipelines-sdk` package. This package has minimal dependencies, but does not
+include first-party TFX components like the TFX ExampleGen, Transform and
+Trainer in `tfx.components.*`, nor any additional tools requiring these
+components.
+
+Both the `tfx` and `ml-pipelines-sdk` packages share the `tfx` namespace and
+the same source repository at https://github.com/tensorflow/tfx. These two
+packages can be built using the instructions below. During development, a
+single editable package may be installed for convenience (see the "Installing
+the development-only `tfx-dev` package" section below).
+
+# Building TFX pip package from source
+
+## Setting up the build environment
+
+First, set up the build environment by running:
+
+```
+package_build/initialize.sh
+```
+
+## Building the `tfx` and `ml-pipelines-sdk` packages
+
+Next, each package can be built using the `bdist_wheel` command:
+
+```
+python package_build/ml-pipelines-sdk/setup.py bdist_wheel
+python package_build/tfx/setup.py bdist_wheel
+```
+
+As a result, `.whl` files will be generated in the `dist/` directory.
+
+# Installing the development-only `tfx-dev` package
+
+During development, it is convenient to install a single editable pip package.
+This package will contain the union of the `tfx` and `ml-pipelines-sdk`
+package in an editable environment. To install this combined package for
+development, run from the repository root:
+
+```
+pip install -e .
+```
+
+This `tfx-dev` package should not be packaged as a binary or source
+distribution using `python setup.py {bdist_wheel,sdist}` to avoid conflicts
+with the two official `tfx` and `ml-pipelines-sdk` packages. Instead, users
+should build the two packages for distribution with the directions above.

package_build/initialize.sh

@@ -0,0 +1,36 @@
+#!/bin/bash
+# Copyright 2020 Google LLC. All Rights Reserved.
+#
+# Licensed 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.
+#
+# Initialization script for building TFX SDK release packages.
+#
+# After this script is run, `python setup.py` commands can be run in the
+# `tfx/` and `ml-pipelines-sdk/` packages.
+
+BASEDIR=$(dirname "$(pwd)/${0#./}")/..
+
+mkdir -p $BASEDIR/dist
+
+for CONFIG_NAME in tfx ml-pipelines-sdk
+do
+  ln -sf $BASEDIR/setup.py $BASEDIR/package_build/$CONFIG_NAME/
+  ln -sf $BASEDIR/dist $BASEDIR/package_build/$CONFIG_NAME/
+  ln -sf $BASEDIR/tfx $BASEDIR/package_build/$CONFIG_NAME/
+  ln -sf $BASEDIR/README*.md $BASEDIR/package_build/$CONFIG_NAME/
+
+  rm -rf $BASEDIR/package_build/$CONFIG_NAME/build
+  mkdir $BASEDIR/package_build/$CONFIG_NAME/build
+  ln -sf $BASEDIR/build/BUILD $BASEDIR/package_build/$CONFIG_NAME/build/
+  ln -sf $BASEDIR/build/gen_proto.sh $BASEDIR/package_build/$CONFIG_NAME/build/
+done

tfx/__init__.py package_build/ml-pipelines-sdk/package_config.py

@@ -1,5 +1,5 @@
-# Lint as: python2, python3
-# Copyright 2019 Google LLC. All Rights Reserved.
+# Lint as: python3
+# Copyright 2020 Google LLC. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,8 +12,8 @@
 # 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.
+"""Configuration for the "ml-pipelines-sdk" package.
 
-"""Init module for TFX."""
-
-# Import version string.
-from tfx.version import __version__
+Core TFX pipeline authoring SDK, with a minimal set of dependencies.
+"""
+PACKAGE_NAME = 'ml-pipelines-sdk'

package_build/tfx/package_config.py

@@ -0,0 +1,21 @@
+# Lint as: python3
+# Copyright 2020 Google LLC. All Rights Reserved.
+#
+# Licensed 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.
+"""Configuration for the "tfx" package.
+
+Recommended installation package for TFX. This package builds on top of
+the "ml-pipelines-sdk" component-authoring SDK package and adds first-party TFX
+components and additional functionality.
+"""
+PACKAGE_NAME = 'tfx'

package_config.py

@@ -0,0 +1,26 @@
+# Lint as: python3
+# Copyright 2020 Google LLC. All Rights Reserved.
+#
+# Licensed 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.
+"""Configuration for the "tfx-dev" package.
+
+Monolithic development package with the entirety of `tfx.*` and the full
+set of dependencies.
+
+Once installed, this is functionally equivalent to the union of the "tfx" and
+"ml-pipeline-sdk" packages, and thus cannot be installed together with the
+latter two packages.
+
+See `package_build/README.md` for packaging details.
+"""
+PACKAGE_NAME = 'tfx-dev'

setup.py

@@ -16,12 +16,13 @@
 
 from __future__ import print_function
 
+import logging
 import os
 import subprocess
 import sys
 
 import setuptools
-from setuptools import find_packages
+from setuptools import find_namespace_packages
 from setuptools import setup
 from setuptools.command import develop
 # pylint: disable=g-bad-import-order
@@ -37,6 +38,15 @@
 from tfx.tools import resolve_deps
 from wheel import bdist_wheel
 
+# Prefer to import `package_config` from the setup.py script's directory. The
+# `package_config.py` file is used to configure which package to build (see
+# the logic below switching on `package_config.PACKAGE_NAME`) and the overall
+# package build README at `package_build/README.md`.
+sys.path.insert(0, os.path.dirname(__file__))
+# pylint: disable=g-bad-import-order,g-import-not-at-top
+import package_config
+# pylint: enable=g-bad-import-order,g-import-not-at-top
+
 
 class _BdistWheelCommand(bdist_wheel.bdist_wheel):
   """Overrided bdist_wheel command.
@@ -160,13 +170,67 @@ def run(self):
         env=os.environ)
 
 
-# Get the long description from the README file.
+# Get the long descriptions from README files.
 with open('README.md') as fp:
-  _LONG_DESCRIPTION = fp.read()
+  _TFX_LONG_DESCRIPTION = fp.read()
+with open('README.ml-pipelines-sdk.md') as fp:
+  _PIPELINES_SDK_LONG_DESCRIPTION = fp.read()
+
+package_name = package_config.PACKAGE_NAME
+tfx_extras_requires = {
+    # In order to use 'docker-image' or 'all', system libraries specified
+    # under 'tfx/tools/docker/Dockerfile' are required
+    'docker-image': dependencies.make_extra_packages_docker_image(),
+    'tfjs': dependencies.make_extra_packages_tfjs(),
+    'examples': dependencies.make_extra_packages_examples(),
+    'test': dependencies.make_extra_packages_test(),
+    'all': dependencies.make_extra_packages_all(),
+}
+ml_pipelines_sdk_packages = ['tfx.dsl', 'tfx.dsl.*']
 
+# This `setup.py` file can be used to build packages in 3 configurations. See
+# the discussion in `package_build/README.md` for an overview. The `tfx` and
+# `ml-pipelines-sdk` pip packages can be built for distribution using the
+# selectable `package_config.PACKAGE_NAME` specifier. Additionally, for
+# development convenience, the `tfx-dev` package containing the union of the
+# the `tfx` and `ml-pipelines-sdk` package can be installed as an editable
+# package using `pip install -e .`, but should not be built for distribution.
+if package_config.PACKAGE_NAME == 'tfx-dev':
+  # Monolithic development package with the entirety of `tfx.*` and the full
+  # set of dependencies. Functionally equivalent to the union of the "tfx" and
+  # "tfx-pipeline-sdk" packages.
+  install_requires = dependencies.make_required_install_packages()
+  extras_require = tfx_extras_requires
+  long_description = _TFX_LONG_DESCRIPTION
+  packages = find_namespace_packages(include=['tfx', 'tfx.*'])
+  # TODO(b/174503231): Remove the override on the following line and prevent
+  # build of wheels from a monolithic "tfx" package.
+  package_name = 'tfx'
+elif package_config.PACKAGE_NAME == 'ml-pipelines-sdk':
+  # Core TFX pipeline authoring SDK, without dependency on component-specific
+  # packages like "tensorflow" and "apache-beam".
+  install_requires = dependencies.make_pipeline_sdk_required_install_packages()
+  extras_require = {}
+  long_description = _PIPELINES_SDK_LONG_DESCRIPTION
+  packages = find_namespace_packages(include=ml_pipelines_sdk_packages)
+elif package_config.PACKAGE_NAME == 'tfx':
+  # Recommended installation package for TFX. This package builds on top of
+  # the "ml-pipelines-sdk" pipeline authoring SDK package and adds first-party
+  # TFX components and additional functionality.
+  install_requires = (
+      ['ml-pipelines-sdk==%s' % version.__version__] +
+      dependencies.make_required_install_packages())
+  extras_require = tfx_extras_requires
+  long_description = _TFX_LONG_DESCRIPTION
+  packages = find_namespace_packages(
+      include=['tfx', 'tfx.*'], exclude=ml_pipelines_sdk_packages)
+else:
+  raise ValueError('Invalid package config: %r.' % package_config.PACKAGE_NAME)
+
+logging.info('Executing build for package %r.', package_name)
 
 setup(
-    name='tfx',
+    name=package_name,
     version=version.__version__,
     author='Google LLC',
     author_email='[email protected]',
@@ -192,16 +256,8 @@ def run(self):
         'Topic :: Software Development :: Libraries :: Python Modules',
     ],
     namespace_packages=[],
-    install_requires=dependencies.make_required_install_packages(),
-    extras_require={
-        # In order to use 'docker-image' or 'all', system libraries specified
-        # under 'tfx/tools/docker/Dockerfile' are required
-        'docker-image': dependencies.make_extra_packages_docker_image(),
-        'tfjs': dependencies.make_extra_packages_tfjs(),
-        'examples': dependencies.make_extra_packages_examples(),
-        'test': dependencies.make_extra_packages_test(),
-        'all': dependencies.make_extra_packages_all(),
-    },
+    install_requires=install_requires,
+    extras_require=extras_require,
     # TODO(b/158761800): Move to [build-system] requires in pyproject.toml.
     setup_requires=[
         'pytest-runner',
@@ -219,10 +275,10 @@ def run(self):
         'resolve_deps': resolve_deps.ResolveDepsCommand,
     },
     python_requires='>=3.6,<3.9',
-    packages=find_packages(),
+    packages=packages,
     include_package_data=True,
     description='TensorFlow Extended (TFX) is a TensorFlow-based general-purpose machine learning platform implemented at Google',
-    long_description=_LONG_DESCRIPTION,
+    long_description=long_description,
     long_description_content_type='text/markdown',
     keywords='tensorflow tfx',
     url='https://www.tensorflow.org/tfx',

tfx/dependencies.py

@@ -50,12 +50,25 @@ def select_constraint(default, nightly=None, git_master=None):
     return default
 
 
+def make_pipeline_sdk_required_install_packages():
+  return [
+      'absl-py>=0.9,<0.11',
+      'ml-metadata' + select_constraint(
+          # LINT.IfChange
+          default='>=0.25,<0.26',
+          # LINT.ThenChange(tfx/workspace.bzl)
+          nightly='>=0.26.0.dev',
+          git_master='@git+https://github.com/google/ml-metadata@master'),
+      'protobuf>=3.12.2,<4',
+      'six>=1.10,<2',
+  ]
+
+
 def make_required_install_packages():
   # Make sure to sync the versions of common dependencies (absl-py, numpy,
   # six, and protobuf) with TF.
   # TODO(b/130767399): add flask once the frontend is exposed externally.
-  return [
-      'absl-py>=0.9,<0.11',
+  return make_pipeline_sdk_required_install_packages() + [
       # LINT.IfChange
       'apache-beam[gcp]>=2.25,<3',
       # LINT.ThenChange(examples/chicago_taxi_pipeline/setup/setup_beam.sh)
@@ -69,16 +82,8 @@ def make_required_install_packages():
       # dependency expecatation with TensorFlow is sorted out.
       'keras-tuner>=1,<1.0.2',
       'kubernetes>=10.0.1,<12',
-      'ml-metadata' + select_constraint(
-          # LINT.IfChange
-          default='>=0.25,<0.26',
-          # LINT.ThenChange(tfx/workspace.bzl)
-          nightly='>=0.26.0.dev',
-          git_master='@git+https://github.com/google/ml-metadata@master'),
-      'protobuf>=3.12.2,<4',
       'pyarrow>=0.17,<0.18',
       'pyyaml>=3.12,<6',
-      'six>=1.10,<2',
       'tensorflow>=1.15.2,!=2.0.*,!=2.1.*,!=2.2.*,<3',
       'tensorflow-hub>=0.9.0,<0.10',
       # TODO(b/159488890): remove user module-only dependency.