Fix tests

- Remove black as docs dependency
- Revert inadvertent scikit-learn version number change
- Remove doc dependencies from `all` optional target
- Fix tfx.v1.proto.__init__ to correctly import the protobufs
- For ci-test.yml, install in normal mode (not editable)

.github/workflows/cd-docs.yml

@@ -39,7 +39,7 @@ jobs:
             mkdocs-material-
 
       - name: Install Dependencies
-        run: pip install mkdocs mkdocs-material mkdocstrings[python] griffe-inherited-docstrings mkdocs-autorefs black mkdocs-jupyter mkdocs-caption
+        run: pip install mkdocs mkdocs-material mkdocstrings[python] griffe-inherited-docstrings mkdocs-autorefs mkdocs-jupyter mkdocs-caption
 
       - name: Deploy to GitHub Pages
         run: mkdocs gh-deploy --force

.github/workflows/ci-test.yml

@@ -52,7 +52,7 @@ jobs:
         python -m pip install --upgrade pip wheel
         # TODO(b/232490018): Cython need to be installed separately to build pycocotools.
         python -m pip install Cython -c ./test_constraints.txt
-        pip install -c ./test_constraints.txt --extra-index-url https://pypi-nightly.tensorflow.org/simple --pre --editable .[all]
+        pip install -c ./test_constraints.txt --extra-index-url https://pypi-nightly.tensorflow.org/simple --pre .[all]
       env:
         TFX_DEPENDENCY_SELECTOR: ${{ matrix.dependency-selector }}
 

tfx/dependencies.py

@@ -98,7 +98,7 @@ def make_required_install_packages():
         # TODO(b/332616741): Scipy version 1.13 breaks the TFX OSS test.
         # Unpin once the issue is resolved.
         "scipy<1.13",
-        "scikit-learn>=1.0,<2",
+        'scikit-learn==1.5.1',
         # TODO(b/291837844): Pinned pyyaml to 5.3.1.
         # Unpin once the issue with installation is resolved.
         "pyyaml>=6,<7",
@@ -267,15 +267,14 @@ def make_extra_packages_docs() -> list[str]:
         "mkdocs-material",
         "griffe-inherited-docstrings",
         "mkdocs-autorefs",
-        "black",
         "mkdocs-jupyter",
         "mkdocs-caption",
         "pymdown-extensions",
     ]
 
 
 def make_extra_packages_all():
-    # All extra dependencies.
+    # All extra dependencies, not including lint or docs dependencies
     return [
         *make_extra_packages_test(),
         *make_extra_packages_tfjs(),
@@ -284,5 +283,4 @@ def make_extra_packages_all():
         *make_extra_packages_tfdf(),
         *make_extra_packages_flax(),
         *make_extra_packages_examples(),
-        *make_extra_packages_docs(),
     ]

tfx/v1/proto/__init__.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 """TFX proto module."""
 
-from tfx.proto import distribution_validator_pb2, example_diff_pb2
 from tfx.proto.bulk_inferrer_pb2 import (
     ClassifyOutput,
     DataSpec,
@@ -24,7 +23,15 @@
     PredictOutputCol,
     RegressOutput,
 )
+from tfx.proto.distribution_validator_pb2 import (
+    DistributionValidatorConfig,
+    FeatureComparator,
+)
 from tfx.proto.evaluator_pb2 import FeatureSlicingSpec, SingleSlicingSpec
+from tfx.proto.example_diff_pb2 import (
+    ExampleDiffConfig,
+    PairedExampleSkew,
+)
 from tfx.proto.example_gen_pb2 import (
     CustomConfig,
     Input,
@@ -46,7 +53,6 @@
     ValidationSpec,
 )
 from tfx.proto.pusher_pb2 import PushDestination, Versioning
-from tfx.proto.pusher_pb2.PushDestination import Filesystem
 from tfx.proto.range_config_pb2 import RangeConfig, RollingRange, StaticRange
 from tfx.proto.trainer_pb2 import EvalArgs, TrainArgs
 from tfx.proto.transform_pb2 import SplitsConfig
@@ -172,7 +178,7 @@
 For example TF Serving only accepts an integer version that is monotonically increasing.
 """
 
-Filesystem.__doc__ = """
+PushDestination.Filesystem.__doc__ = """
 File system based destination definition.
 """
 
@@ -212,26 +218,18 @@
 Args specific to tuning in `components.Tuner`.
 """
 
-ExampleDiffConfig = example_diff_pb2.ExampleDiffConfig
-
 ExampleDiffConfig.__doc__ = """
 Configurations related to Example Diff.
 """
 
-FeatureComparator = distribution_validator_pb2.FeatureComparator
-
 FeatureComparator.__doc__ = """
 Per feature configuration in Distribution Validator.
 """
 
-DistributionValidatorConfig = distribution_validator_pb2.DistributionValidatorConfig
-
 DistributionValidatorConfig.__doc__ = """
 Configurations related to Distribution Validator.
 """
 
-PairedExampleSkew = example_diff_pb2.PairedExampleSkew
-
 PairedExampleSkew.__doc__ = """
 Configurations related to Example Diff on feature pairing level.
 """