Cannot customize resolver strategy after TFX 1.7.0

System information

• Have I specified the code to reproduce the issue (Yes, No): Yes
• Environment in which the code is executed (e.g., Local(Linux/MacOS/Windows), Interactive Notebook, Google Cloud, etc): Local(MacOS, Linux)
• TensorFlow version: tensorflow==2.8.0
• TFX Version: tfx==1.7.0
• Python version: Python 3.8.7
• Python dependencies (from pip freeze output):

Describe the current behavior

First of all, I’m aware that the tfx.dsl.input_resolution.strategies module is experimental.

Before TFX 1.7, it was able to run a custom resolver strategy using LocalDagRunner, but after version 1.7, I cannot run my own strategy.

Describe the expected behavior

Resolve can be run with custom resolver strategy that subclasses tfx.dsl.components.common.resolver.ResolverStrategy.

Standalone code to reproduce the issue

I wrote the sample code that runs with custom strategy and prints execution results.

import json
from typing import Dict, List, Optional

import ml_metadata as mlmd
from ml_metadata.proto import metadata_store_pb2
from tfx import types
from tfx.dsl.components.common import resolver
from tfx.orchestration import metadata
from tfx.v1 import dsl
from tfx.v1.orchestration import LocalDagRunner
from tfx.v1.types import standard_artifacts


class CustomStrategy(resolver.ResolverStrategy):
    def __init__(self, desired_num_of_artifacts: int = 1):
        pass

    def resolve_artifacts(
        self, store: mlmd.MetadataStore, input_dict: Dict[str, List[types.Artifact]]
    ) -> Optional[Dict[str, List[types.Artifact]]]:
        result = {}
        for k, artifact_list in input_dict.items():
            result[k] = artifact_list[1:2]
        return result


custom_resolver = dsl.Resolver(
    strategy_class=CustomStrategy,
    statistics=dsl.Channel(type=standard_artifacts.ExampleStatistics),
    config={"desired_num_of_artifacts": 1},
).with_id("custom_resolver")
connection_config = metadata_store_pb2.ConnectionConfig()
connection_config.sqlite.filename_uri = "./tmp/metadata-sqlite"
resolver_pipeline = dsl.Pipeline(
    pipeline_name="Hello",
    pipeline_root="./tmp",
    metadata_connection_config=connection_config,
    components=[custom_resolver],
)
LocalDagRunner().run(resolver_pipeline)

with metadata.Metadata(connection_config=connection_config) as m:
    execution_result = m.store.get_executions()[0].custom_properties["__execution_result__"].string_value
    print(json.loads(execution_result)["resultMessage"])

execution result:

WARNING:absl:Artifact type ExampleStatistics is not found in MLMD.
Traceback (most recent call last):

  File "SITE_PACKAGES_PATH/tfx/orchestration/portable/inputs_utils.py", line 243, in resolve_input_artifacts_v2
    resolved = processor.run_resolver_steps(

  File "SITE_PACKAGES_PATH/tfx/orchestration/portable/input_resolution/processor.py", line 87, in run_resolver_steps
    cls = ops.get_by_class_path(step.class_path)

  File "SITE_PACKAGES_PATH/tfx/dsl/input_resolution/ops/ops.py", line 52, in get_by_class_path
    return _OPS_BY_CLASSPATH[class_path]

KeyError: '__main__.CustomStrategy'


The above exception was the direct cause of the following exception:


Traceback (most recent call last):

  File "SITE_PACKAGES_PATH/tfx/orchestration/portable/resolver_node_handler.py", line 74, in run
    resolved_inputs = inputs_utils.resolve_input_artifacts_v2(

  File "SITE_PACKAGES_PATH/tfx/orchestration/portable/inputs_utils.py", line 252, in resolve_input_artifacts_v2
    raise exceptions.InputResolutionError(

tfx.orchestration.portable.input_resolution.exceptions.InputResolutionError: Error occurred during input resolution: '__main__.CustomStrategy'.

Name of your Organization (Optional)

Other info / logs

Since resolver strategy class object is fetched by using tfx.dsl.input_resolution.ops.ops.get_by_class_path in tfx 1.7, I checked that it can run with tfx.dsl.input_resolution.ops.ops.register function in LocalDagRunner, but it can’t run in KubeFlowDagRunner. (I couldn’t modify the entrypoint of kubeflow dag runner)

–

Thanks 😃

[chongkong]

TFX 1.7.1 is also released including the patch! thanks

[jeongukjae]

@chongkong Thanks 😃 but just one thing to say is that my name is not Jeonguk, it’s Ukjae (욱재) (It’s because of the name order in Korea 😅)