switch between the 2 potential types

Author: casassg

tfx/orchestration/local/local_dag_runner.py

@@ -50,7 +50,10 @@ def run_with_ir(self, pipeline: pipeline_pb2.Pipeline) -> None:
         })
 
     deployment_config = runner_utils.extract_local_deployment_config(pipeline)
-    connection_config = deployment_config.metadata_connection_config or deployment_config.mlmd_store_client_config
+    if deployment_config.HasField('metadata_connection_config'):
+      connection_config = deployment_config.metadata_connection_config
+    else:
+      connection_config = deployment_config.mlmd_store_client_config
 
     logging.info('Using deployment config:\n %s', deployment_config)
     logging.info('Using connection config:\n %s', connection_config)

tfx/orchestration/local/runner_utils.py

@@ -92,8 +92,8 @@ def _to_local_deployment(
         _build_local_platform_config(k, v))
 
   if not input_config.metadata_connection_config.Unpack(
-      result.metadata_connection_config) or not input_config.mlmd_store_client_config.Unpack(
-        result.metadata_connection_config):
+      result.metadata_connection_config) and not input_config.metadata_connection_config.Unpack(
+        result.mlmd_store_client_config):
     raise ValueError('metadata_connection_config is expected to be in type '
                      'ml_metadata.ConnectionConfig or ml_metadata.MetadataStoreClientConfig, but got type {}'.format(
                          input_config.metadata_connection_config.type_url))

tfx/proto/orchestration/local_deployment_config.proto

@@ -47,11 +47,11 @@ message LocalDeploymentConfig {
   // set custom driver logic.
   map<string, ExecutableSpec> custom_driver_specs = 2;
   // Connection config for ML-metadata.
-  optional ml_metadata.ConnectionConfig metadata_connection_config = 3;
+  ml_metadata.ConnectionConfig metadata_connection_config = 3;
   // A key from `node_id` to platform configs. Specifies platform settings to
   // pass to individual node operators.
   map<string, LocalPlatformConfig> node_level_platform_configs = 4;
   // Server Connection config
-  optional ml_metadata.MetadataStoreClientConfig mlmd_store_client_config = 5;
+  ml_metadata.MetadataStoreClientConfig mlmd_store_client_config = 5;
 
 }