[Upstream] [COR-2297/] Fix nested offloaded type validation (#552)

[pmahindrakar-oss]

The following workflow works when we are not offloading literals in flytekit

import logging
from typing import List
from flytekit import map_task, task, workflow,LaunchPlan

logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger("flytekit")
logger.setLevel(logging.DEBUG)

@task(cache=True, cache_version="1.1")
def my_30mb_task(i: str) -> str:
    return f"Hello world {i}" * 30 * 100 * 1024

@task(cache=True, cache_version="1.1")
def generate_strs(count: int) -> List[str]:
    return ["a"] * count

@workflow
def my_30mb_wf(mbs: int) -> List[str]:
  strs = generate_strs(count=mbs)
  return map_task(my_30mb_task)(i=strs)

@workflow
def big_inputs_wf(input: List[str]):
   noop()

@task(cache=True, cache_version="1.1")
def noop():
    ...

big_inputs_wf_lp = LaunchPlan.get_or_create(name="big_inputs_wf_lp", workflow=big_inputs_wf)

@workflow
def ref_wf(mbs: int):
  big_inputs_wf_lp(input=my_30mb_wf(mbs))

Without flytekit offloading the return type is OffloadedLiteral{inferredType:{Collection{String}} and when checked against big_inputs_wf launchplan which needs Collection{String} , the LiteralTypeToLiteral returns the inferredType : Collection{String}

If we enable offloading in flytekit, the returned data from map task is Collection{OffloadedLiteral<{inferredType:{Collection{String}}}

When passing this Input to big_inputs_wf which takes Collection{String} then the type check fails due to LiteralTypeToLiteral returning Collection{OffloadedLiteral{inferredType:{Collection{String}}} as Collection{Collection{String}}

Flytekit handles this case by special casing Collection{OffloadedLiteral} and similar special casing is needed in flyte code base

Tested this by deploying this PR changes

https://dogfood-gcp.cloud-staging.union.ai/console/projects/flytesnacks/domains/development/executions/akxs97cdmkmxhhqp228x/nodes

Earlier it would fail like this https://dogfood-gcp.cloud-staging.union.ai/console/projects/flytesnacks/domains/development/executions/ap4thjp5528kjfspcsds/nodes

[UserError] failed to launch workflow, caused by: rpc error: code = InvalidArgument desc = invalid input input wrong type. Expected collection_type:{simple:STRING}, but got collection_type:{collection_type:{simple:STRING}}

Rollout to canary and then all prod byoc and serverless tenants

Should this change be upstreamed to OSS (flyteorg/flyte)? If not, please uncheck this box, which is used for auditing. Note, it is the responsibility of each developer to actually upstream their changes. See this guide.

• To be upstreamed to OSS

TODO: Link Linear issue(s) using magic words. fixes will move to merged status, while ref will only link the PR.

• Added tests
• Ran a deploy dry run and shared the terraform plan
• Added logging and metrics
• Updated dashboards and alerts
• Updated documentation

Tracking issue

Why are the changes needed?

What changes were proposed in this pull request?

How was this patch tested?

Setup process

Screenshots

Check all the applicable boxes

• I updated the documentation accordingly.
• All new and existing tests passed.
• All commits are signed-off.

Related PRs

Docs link

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 36.95%. Comparing base (3c3ae05) to head (7f0db4a).
Report is 1 commits behind head on master.

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #5996   +/-   ##
=======================================
  Coverage   36.95%   36.95%           
=======================================
  Files        1310     1310           
  Lines      131451   131464   +13     
=======================================
+ Hits        48575    48586   +11     
- Misses      78655    78657    +2     
  Partials     4221     4221           

Flag	Coverage Δ
unittests-datacatalog	51.58% <ø> (ø)
unittests-flyteadmin	54.05% <ø> (-0.03%)	⬇️
unittests-flytecopilot	22.23% <ø> (ø)
unittests-flytectl	62.44% <ø> (+0.04%)	⬆️
unittests-flyteidl	6.95% <ø> (ø)
unittests-flyteplugins	53.83% <ø> (ø)
unittests-flytepropeller	43.10% <100.00%> (+0.02%)	⬆️
unittests-flytestdlib	55.31% <ø> (ø)

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.