Sklearn type transformer should be automatically loaded with import flytekit

.readthedocs.yml

@@ -18,3 +18,4 @@ sphinx:
 python:
   install:
     - requirements: doc-requirements.txt
+    - requirements: docs/requirements.txt

docs/requirements.txt

@@ -0,0 +1,7 @@
+# TODO: Remove after buf migration is done and packages updated, see doc-requirements.in
+# skl2onnx and tf2onnx added here so that the plugin API reference is rendered,
+# with the caveat that the docs build environment has the environment variable
+# PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python set so that protobuf can be parsed
+# using Python, which is acceptable for docs building.
+skl2onnx
+tf2onnx

flytekit/__init__.py

@@ -184,7 +184,7 @@
 from flytekit.core.workflow import ImperativeWorkflow as Workflow
 from flytekit.core.workflow import WorkflowFailurePolicy, reference_workflow, workflow
 from flytekit.deck import Deck
-from flytekit.extras import pytorch, tensorflow
+from flytekit.extras import pytorch, sklearn, tensorflow
 from flytekit.extras.persistence import GCSPersistence, HttpPersistence, S3Persistence
 from flytekit.loggers import logger
 from flytekit.models.common import Annotations, AuthRole, Labels

flytekit/extras/pytorch/__init__.py

@@ -1,13 +1,14 @@
 """
-Flytekit PyTorch
-=========================================
 .. currentmodule:: flytekit.extras.pytorch
 
 .. autosummary::
    :template: custom.rst
    :toctree: generated/
 
     PyTorchCheckpoint
+    PyTorchCheckpointTransformer
+    PyTorchModuleTransformer
+    PyTorchTensorTransformer
 """
 from flytekit.loggers import logger
 

flytekit/extras/sklearn/__init__.py

@@ -1,11 +1,11 @@
 """
-Flytekit Sklearn
-=========================================
 .. currentmodule:: flytekit.extras.sklearn
 
 .. autosummary::
    :template: custom.rst
    :toctree: generated/
+
+   SklearnEstimatorTransformer
 """
 from flytekit.loggers import logger
 

flytekit/extras/tensorflow/__init__.py

@@ -1,3 +1,14 @@
+"""
+.. currentmodule:: flytekit.extras.tensorflow
+
+.. autosummary::
+   :template: custom.rst
+   :toctree: generated/
+
+   TensorFlowRecordFileTransformer
+   TensorFlowRecordsDirTransformer
+"""
+
 from flytekit.loggers import logger
 
 # TODO: abstract this out so that there's an established pattern for registering plugins