\nRe-sync comments\n

PiperOrigin-RevId: 319467552

Author: tfx-copybara

setup.py

@@ -149,4 +149,3 @@ def generate_proto(source):
         [console_scripts]
         tfx=tfx.tools.cli.cli_main:cli_group
     """)
-

tfx/components/base/base_driver_test.py

@@ -106,8 +106,7 @@ def setUp(self):
         pipeline_info=self._pipeline_info)
 
   @mock.patch(
-      'tfx.components.base.base_driver.BaseDriver.verify_input_artifacts'
-  )
+      'tfx.components.base.base_driver.BaseDriver.verify_input_artifacts')
   @mock.patch.object(types.ValueArtifact, 'read', fake_read)
   def testPreExecutionNewExecution(self, mock_verify_input_artifacts_fn):
     self._mock_metadata.search_artifacts.return_value = list(
@@ -142,8 +141,7 @@ def testPreExecutionNewExecution(self, mock_verify_input_artifacts_fn):
                      _STRING_VALUE)
 
   @mock.patch(
-      'tfx.components.base.base_driver.BaseDriver.verify_input_artifacts'
-  )
+      'tfx.components.base.base_driver.BaseDriver.verify_input_artifacts')
   @mock.patch.object(types.ValueArtifact, 'read', fake_read)
   def testPreExecutionCached(self, mock_verify_input_artifacts_fn):
     self._mock_metadata.search_artifacts.return_value = list(

tfx/components/transform/executor_test.py

@@ -183,7 +183,7 @@ def testCounters(self):
 
     # The test data has 10036 instances in the train dataset, and 4964 instances
     # in the eval dataset (obtained by running:
-    #   gqui third_party/tfx/components/testdata/csv_example_gen/train/data* \
+    #   gqui third_party/py/tfx/components/testdata/csv_example_gen/train/data* \
     #     'select count(*)'
     # )
     # Since the analysis dataset (train) is read twice (once for analysis and

tfx/components/tuner/executor.py

@@ -24,14 +24,14 @@
 from absl import logging
 from kerastuner.engine import base_tuner
 
-from google.protobuf import json_format
 from tfx import types
 from tfx.components.base import base_executor
 from tfx.components.trainer import fn_args_utils
 from tfx.components.util import udf_utils
 from tfx.proto import tuner_pb2
 from tfx.types import artifact_utils
 from tfx.utils import io_utils
+from google.protobuf import json_format
 
 # Key for best hyperparameters in executor output_dict.
 _BEST_HYPERPARAMETERS_KEY = 'best_hyperparameters'

tfx/dsl/compiler/testdata/iris_pipeline_ir.pbtxt

@@ -2,7 +2,7 @@
 # proto-message: Pipeline
 #
 # This file contains the IR of an adapted Iris example pipeline
-# third_party/tfx/examples/iris/iris_pipeline_beam.py
+# tfx/examples/iris/iris_pipeline_beam.py
 
 pipeline_info {
   id: "iris"

tfx/examples/custom_components/presto_example_gen/proto/presto_config.proto

@@ -18,26 +18,42 @@ package tfx.examples.custom_components.presto_example_gen;
 message PrestoConnConfig {
   // Based off official prestodb-python-client library
   // https://github.com/prestodb/presto-python-client
-  string host = 1; // required
-  oneof opt_port { uint32 port = 2; }
-  oneof opt_user { string user = 3; }
-  oneof opt_source { string source = 4; }
-  oneof opt_catalog { string catalog = 5; }
-  oneof opt_schema { string schema = 6; }
+  string host = 1;  // required
+  oneof opt_port {
+    uint32 port = 2;
+  }
+  oneof opt_user {
+    string user = 3;
+  }
+  oneof opt_source {
+    string source = 4;
+  }
+  oneof opt_catalog {
+    string catalog = 5;
+  }
+  oneof opt_schema {
+    string schema = 6;
+  }
   // TODO session_properties = 7
   // TODO http_headers = 8
-  oneof opt_http_scheme { string http_scheme = 9; }
+  oneof opt_http_scheme {
+    string http_scheme = 9;
+  }
   oneof opt_auth {
-    BasicAuthConfig basic_auth = 10; // prestodb.auth.BasicAuthentication class
+    BasicAuthConfig basic_auth = 10;  // prestodb.auth.BasicAuthentication class
     // TODO KerberosAuthConfig kerberos_auth = 11
   }
   // TODO redirect_handler = 12
-  oneof opt_max_attempts { uint32 max_attempts = 13; }
-  oneof opt_request_timeout { uint32 request_timeout = 14; }
+  oneof opt_max_attempts {
+    uint32 max_attempts = 13;
+  }
+  oneof opt_request_timeout {
+    uint32 request_timeout = 14;
+  }
   // TODO isolation_level = 15
 
   message BasicAuthConfig {
-    string username = 1; // required
-    string password = 2; // required
+    string username = 1;  // required
+    string password = 2;  // required
   }
 }

tfx/experimental/templates/taxi/e2e_tests/kubeflow_e2e_test.py

@@ -23,13 +23,12 @@
 from absl import logging
 import kfp
 import tensorflow as tf
-import yaml
-
-from google.cloud import storage
 from tfx.experimental.templates.taxi.e2e_tests import test_utils
 from tfx.orchestration import test_utils as orchestration_test_utils
 from tfx.orchestration.kubeflow import test_utils as kubeflow_test_utils
 from tfx.utils import telemetry_utils
+import yaml
+from google.cloud import storage
 
 
 class TaxiTemplateKubeflowE2ETest(test_utils.BaseEndToEndTest):

tfx/extensions/google_cloud_ai_platform/runner_test.py

@@ -79,9 +79,7 @@ def _serialize_custom_config_under_test(self) -> Dict[Text, Any]:
     result['custom_config'] = json_utils.dumps(result['custom_config'])
     return result
 
-  @mock.patch(
-      'tfx.extensions.google_cloud_ai_platform.runner.discovery'
-  )
+  @mock.patch('tfx.extensions.google_cloud_ai_platform.runner.discovery')
   def testStartAIPTraining(self, mock_discovery):
     mock_discovery.build.return_value = self._mock_api_client
     self._setUpTrainingMocks()
@@ -114,9 +112,7 @@ def testStartAIPTraining(self, mock_discovery):
     self.assertStartsWith(body['jobId'], 'tfx_')
     self._mock_get.execute.assert_called_with()
 
-  @mock.patch(
-      'tfx.extensions.google_cloud_ai_platform.runner.discovery'
-  )
+  @mock.patch('tfx.extensions.google_cloud_ai_platform.runner.discovery')
   def testStartAIPTrainingWithUserContainer(self, mock_discovery):
     mock_discovery.build.return_value = self._mock_api_client
     self._setUpTrainingMocks()
@@ -238,9 +234,7 @@ def _assertDeployModelMockCalls(self,
             self._project_id, self._model_name, self._model_version))
     self._mock_set_default_execute.assert_called_with()
 
-  @mock.patch(
-      'tfx.extensions.google_cloud_ai_platform.runner.discovery'
-  )
+  @mock.patch('tfx.extensions.google_cloud_ai_platform.runner.discovery')
   def testDeployModelForAIPPrediction(self, mock_discovery):
     mock_discovery.build.return_value = self._mock_api_client
     self._setUpPredictionMocks()
@@ -257,9 +251,7 @@ def testDeployModelForAIPPrediction(self, mock_discovery):
     self._assertDeployModelMockCalls(
         expected_models_create_body=expected_models_create_body)
 
-  @mock.patch(
-      'tfx.extensions.google_cloud_ai_platform.runner.discovery'
-  )
+  @mock.patch('tfx.extensions.google_cloud_ai_platform.runner.discovery')
   def testDeployModelForAIPPredictionError(self, mock_discovery):
     mock_discovery.build.return_value = self._mock_api_client
     self._setUpPredictionMocks()
@@ -286,9 +278,7 @@ def testDeployModelForAIPPredictionError(self, mock_discovery):
         expected_models_create_body=expected_models_create_body,
         expect_set_default=False)
 
-  @mock.patch(
-      'tfx.extensions.google_cloud_ai_platform.runner.discovery'
-  )
+  @mock.patch('tfx.extensions.google_cloud_ai_platform.runner.discovery')
   def testDeployModelForAIPPredictionWithCustomRegion(self, mock_discovery):
     mock_discovery.build.return_value = self._mock_api_client
     self._setUpPredictionMocks()
@@ -306,9 +296,7 @@ def testDeployModelForAIPPredictionWithCustomRegion(self, mock_discovery):
     self._assertDeployModelMockCalls(
         expected_models_create_body=expected_models_create_body)
 
-  @mock.patch(
-      'tfx.extensions.google_cloud_ai_platform.runner.discovery'
-  )
+  @mock.patch('tfx.extensions.google_cloud_ai_platform.runner.discovery')
   def testDeployModelForAIPPredictionWithCustomRuntime(self, mock_discovery):
     mock_discovery.build.return_value = self._mock_api_client
     self._setUpPredictionMocks()

tfx/extensions/google_cloud_ai_platform/tuner/executor_test.py

@@ -22,11 +22,11 @@
 # TODO(b/160252429): Clearn up lint warining.
 import tensorflow as tf
 
-from google.protobuf import json_format
 from tfx.extensions.google_cloud_ai_platform.trainer import executor as ai_platform_trainer_executor
 from tfx.extensions.google_cloud_ai_platform.tuner import executor as ai_platform_tuner_executor
 from tfx.proto import tuner_pb2
 from tfx.utils import json_utils
+from google.protobuf import json_format
 
 
 class ExecutorTest(tf.test.TestCase):
@@ -57,8 +57,8 @@ def setUp(self):
 
     self.addCleanup(mock.patch.stopall)
     self.mock_runner = mock.patch(
-        'tfx.extensions.google_cloud_ai_platform.tuner.executor.runner'
-    ).start()
+        'tfx.extensions.google_cloud_ai_platform.tuner.executor.runner').start(
+        )
 
   def _serialize_custom_config_under_test(self) -> Dict[Text, Any]:
     """Converts self._exec_properties['custom_config'] to string."""

tfx/orchestration/airflow/airflow_dag_runner_test.py

@@ -106,9 +106,7 @@ def __init__(self, spec: types.ComponentSpec):
 
 class AirflowDagRunnerTest(tf.test.TestCase):
 
-  @mock.patch(
-      'tfx.orchestration.airflow.airflow_component.AirflowComponent'
-  )
+  @mock.patch('tfx.orchestration.airflow.airflow_component.AirflowComponent')
   @mock.patch('airflow.models.DAG')
   def testAirflowDagRunner(self, mock_airflow_dag_class,
                            mock_airflow_component_class):

tfx/orchestration/experimental/interactive/interactive_context_test.py

@@ -203,9 +203,8 @@ def testExportToPipelineRaisesErrorInvalidRunnerType(self):
                            export_filepath=export_filepath,
                            runner_type='foobar')
 
-  @mock.patch(
-      'tfx.orchestration.experimental.interactive.'
-      'standard_visualizations.ExampleAnomaliesVisualization.display')
+  @mock.patch('tfx.orchestration.experimental.interactive.'
+              'standard_visualizations.ExampleAnomaliesVisualization.display')
   def testShow(self, *unused_mocks):
     context = interactive_context.InteractiveContext()
     mock_object = mock.MagicMock()

tfx/orchestration/experimental/interactive/visualizations_test.py

@@ -37,9 +37,8 @@ def tearDown(self):
     del builtins.__dict__['__IPYTHON__']
     super(VisualizationsTest, self).tearDown()
 
-  @mock.patch(
-      'tfx.orchestration.experimental.interactive.'
-      'visualizations.get_registry')
+  @mock.patch('tfx.orchestration.experimental.interactive.'
+              'visualizations.get_registry')
   def testVisualizationRegistrationAndUsage(self, *unused_mocks):
     registry = visualizations.ArtifactVisualizationRegistry()
     visualizations.get_registry = mock.MagicMock(return_value=registry)

tfx/proto/orchestration/pipeline.proto

@@ -139,7 +139,7 @@ message PropertyPredicate {
     // The left-hand side element to the logical operator.
     PropertyPredicate lhs = 2;
     // The right-hand side element to the logical operator.
-    PropertyPredicate rhs= 3;
+    PropertyPredicate rhs = 3;
   }
   oneof operator {
     ValueComparator value_comparator = 1;

tfx/proto/pusher.proto

@@ -52,4 +52,3 @@ enum Versioning {
   // model is the same.
   UNIX_TIMESTAMP = 1;
 }
-

tfx/proto/trainer.proto

@@ -17,7 +17,6 @@ package tfx.components.trainer;
 
 // Args specific to training.
 message TrainArgs {
-
   // The number of steps to train on.
   int32 num_steps = 2;
 
@@ -26,7 +25,6 @@ message TrainArgs {
 
 // Args specific to eval.
 message EvalArgs {
-
   // The number of evaluation steps to take.
   int32 num_steps = 2;
 

tfx/tools/build_docs.py

@@ -48,7 +48,7 @@
 from tensorflow_docs.api_generator import generate_lib
 
 
-# Standard Imports.third_party.tfx as tfx
+import tfx as tfx
 # pylint: disable=unused-import
 from tfx import components
 from tfx import orchestration

tfx/tools/cli/handler/template_handler.py

@@ -164,9 +164,7 @@ def copy_template(flags_dict: Dict[Text, Any]) -> None:
       _IMPORT_FROM_PACKAGE: _IMPORT_FROM_LOCAL_DIR,
       _PLACEHOLDER_PIPELINE_NAME: pipeline_name,
   }
-  _copy_and_replace_placeholder_dir(template_dir,
-                                    destination_dir,
-                                    ignore_paths,
+  _copy_and_replace_placeholder_dir(template_dir, destination_dir, ignore_paths,
                                     replace_dict)
   for additional_file in _ADDITIONAL_FILE_PATHS.get(model, []):
     src_path = os.path.join(_tfx_src_dir(), additional_file.src)