Repository: kubeflow/pipelines

Issue:
[components]: TrainingPipeline Custom Task accepting multiple datasets defaulting to incorrect dataset type
Originally raised in [stackoverflow - #68025813](https://stackoverflow.com/questions/67999255/error-tabular-dataset-treated-as-image-dataset-vertex-ai-pipelines-custom-tra/68025813#68025813)

## Environment
python 3.7.3
kfp==1.6.2
kfp-pipeline-spec==0.1.7
kfp-server-api==1.6.0
google-cloud-pipeline-components==0.1.1

## Error message
```
ValueError: ImageDataset class can not be used to retrieve dataset resource projects/nnnnnnnnnnnn/locations/us-central1/datasets/3781554739456507904, check the dataset type
    f"{self.__class__.__name__} class can not be used to retrieve "
  File "/opt/python3.7/lib/python3.7/site-packages/google/cloud/aiplatform/datasets/dataset.py", line 100, in _validate_metadata_schema_uri
    self._validate_metadata_schema_uri()
  File "/opt/python3.7/lib/python3.7/site-packages/google/cloud/aiplatform/datasets/dataset.py", line 82, in __init__
    return annotation_type(value)
  File "/opt/python3.7/lib/python3.7/site-packages/google_cloud_pipeline_components/aiplatform/remote_runner.py", line 176, in cast
    value = cast(value, param_type)
  File "/opt/python3.7/lib/python3.7/site-packages/google_cloud_pipeline_components/aiplatform/remote_runner.py", line 205, in prepare_parameters
    prepare_parameters(serialized_args[METHOD_KEY], method, is_init=False)
  File "/opt/python3.7/lib/python3.7/site-packages/google_cloud_pipeline_components/aiplatform/remote_runner.py", line 236, in runner
    print(runner(args.cls_name, args.method_name, executor_input, kwargs))
  File "/opt/python3.7/lib/python3.7/site-packages/google_cloud_pipeline_components/aiplatform/remote_runner.py", line 280, in main
    main()
  File "/opt/python3.7/lib/python3.7/site-packages/google_cloud_pipeline_components/aiplatform/remote_runner.py", line 284, in <module>
    exec(code, run_globals)
  File "/opt/python3.7/lib/python3.7/runpy.py", line 85, in _run_code
    "__main__", mod_spec)
  File "/opt/python3.7/lib/python3.7/runpy.py", line 193, in _run_module_as_main
Traceback (most recent call last):
```

## Sample pipeline:
```python
import datetime
from kfp.v2 import dsl, compiler
from kfp.v2.google.client import AIPlatformClient
import google_cloud_pipeline_components.aiplatform as gcc_ai

PROJECT = "my-project"
PIPELINE_NAME = "test-pipeline"
PIPELINE_ROOT_PATH = f"gs://test-pipeline-20210525/{PIPELINE_NAME}"

@dsl.pipeline(
  name=PIPELINE_NAME,
  pipeline_root=PIPELINE_ROOT_PATH
)
def test_pipeline(
  display_name: str=f"{PIPELINE_NAME}-2021MMDD-nn"
):
  dataset_create_op = gcc_ai.TabularDatasetCreateOp(
    project=PROJECT, display_name=display_name,
    gcs_source="gs://used_apartment/datasource/train.csv"
  )

  training_job_run_op = gcc_ai.CustomContainerTrainingJobRunOp(
    project=PROJECT, display_name=display_name,
    container_uri="us-central1-docker.pkg.dev/my-project/dataops-rc2021/custom-train:latest",
    staging_bucket="vertex_ai_staging_rc2021",
    base_output_dir="gs://used_apartment/cstm_img_scrf/artifact",
    model_serving_container_image_uri="us-central1-docker.pkg.dev/my-project/dataops-rc2021/custom-pred:latest",
    model_serving_container_predict_route="/",
    model_serving_container_health_route="/health",
    model_serving_container_ports=[8080],
    training_fraction_split=0.8,
    validation_fraction_split=0.1,
    test_fraction_split=0.1,
    dataset=dataset_create_op.outputs["dataset"]
  )
```


Discussion:

Resolving patch:
diff --git a/components/google-cloud/google_cloud_pipeline_components/aiplatform/utils.py b/components/google-cloud/google_cloud_pipeline_components/aiplatform/utils.py
--- a/components/google-cloud/google_cloud_pipeline_components/aiplatform/utils.py
+++ b/components/google-cloud/google_cloud_pipeline_components/aiplatform/utils.py
@@ -69,6 +69,11 @@ def get_aiplatform_class_by_name(_annotation):
         pass
 
 
+# This is the Union of all typed datasets.
+# Relying on the annotation defined in the SDK
+# as additional typed Datasets may be added in the future.
+dataset_annotation = inspect.signature(aiplatform.CustomTrainingJob.run).parameters['dataset'].annotation
+
 def resolve_annotation(annotation: Any) -> Any:
     """Resolves annotation type against a MB SDK type.
 
@@ -87,6 +92,11 @@ def resolve_annotation(annotation: Any) -> Any:
         if issubclass(annotation, aiplatform.base.VertexAiResourceNoun):
             return annotation
 
+    # if this is a union of all typed datasets annotation
+    if annotation is dataset_annotation:
+        # return the catch all Dataset class
+        return aiplatform.datasets.dataset._Dataset
+
     # handle forward references
     resolved_annotation = get_forward_reference(annotation)
     if resolved_annotation:


Test patch:
diff --git a/components/google-cloud/tests/aiplatform/test_utils.py b/components/google-cloud/tests/aiplatform/test_utils.py
--- a/components/google-cloud/tests/aiplatform/test_utils.py
+++ b/components/google-cloud/tests/aiplatform/test_utils.py
@@ -248,3 +248,17 @@ def test_generate_docstring_with_PARAMS_TO_REMOVE(self):
             "        input_param\n"
         )
         self.assertEqual(results, expected_results)
+
+    def test_custom_training_typed_dataset_annotation_defaults_to_using_base_dataset(self):
+
+        dataset_annotation = signature(aiplatform.CustomTrainingJob.run).parameters['dataset'].annotation
+
+        assert utils.resolve_annotation(dataset_annotation) is aiplatform.datasets.dataset._Dataset
+
+        dataset_annotation = signature(aiplatform.CustomContainerTrainingJob.run).parameters['dataset'].annotation
+
+        assert utils.resolve_annotation(dataset_annotation) is aiplatform.datasets.dataset._Dataset
+
+        dataset_annotation = signature(aiplatform.CustomPythonPackageTrainingJob.run).parameters['dataset'].annotation
+
+        assert utils.resolve_annotation(dataset_annotation) is aiplatform.datasets.dataset._Dataset