[UPath I/O managers] special case handling of None outputs

[jamiedemaria]

Summary & Motivation

Long term, we’d like to have the following system re: I/O managers

• Behavior that opts you into using the I/O manager:
  ** returning any non-None value from an asset or op
  ** doing a parameter-based dependency (def my_asset(upstream))
• Behavior that does not use the I/O manager:
  ** returning None from an asset or op
  ** therefore you should set up dependencies using deps (Nothing dependencies for op).
  ** If you do a parameter-based dependency on an asset that returns None, the expectation is that you have met the contract of handle_output in some other way

The DB IO managers already error when None is returned, but the UPath I/O manager does store Nones, which means that this setup

@asset 
def upstream():
    return None

@asset 
def downstream(upstream):
    ....

will break if we go directly to not storing all returned Nones. However, creating a lot of files that just store None causes other problems:

1. Vast number of vestigial files created in cases when the asset is not being loaded as an I/O-managed input. This is especially apparent for large backfills.
2. Moving an asset from non-partitioned to partitioned can cause errors when the partitioned file path cannot exist

So as a workaround we will add special behavior to the UPath I/O manager so that None values will not be stored in the file system, but will still be loadable by the I/O manager in downstream assets. We will not add special behavior for ops because it conflicts with expectations for memoization. Once memoization is deprecated and removed we can add special casing for ops based on whether the expected output file exists. If it does not exist we can assume the output was None and provide None.

For assets:
In execute_step.py we will add a "marker" when a None value is handled. In UPath IO manager handle_output we will not store the None in the file system. Then in load_input, if the corresponding output has the "marker" we will automatically return None rather that reading from the file system. The "marker" will be stored as tags, which are surfaced as part of the data versioning code in #19324 so this does not add any additional reads to the db

Making these changes will allow us to continue supporting None values with the UPath I/O manager with no change in behavior, but also solve issues 1 and 2 described above.

To Do:

• ensure logging messages about outputs are updated to accurately reflect the changes in this PR. For example, we should not log a path to a file when no file is created to store None

Notes:
reference PR #15611

How I Tested These Changes

[jamiedemaria]

Warning

This pull request is not mergeable via GitHub because a downstack PR is open. Once all requirements are satisfied, merge this PR as a stack on Graphite.
Learn more

Current dependencies on/for this PR:

• master
  • PR Only have one kind of context for direct invocation #17554
    • PR duplicate op context methods on asset context  #18573
      • PR deprecate methods on AssetExecutionContext that can be found via the DagsterRun object #17339
        • PR Direct invocation asset execution context #18549
          • PR deprecate asset_partition_*_for_output on AssetExecutionContext #19436
            • PR [docs] remove asset_partition_*_for_output from docs #19437
              • PR deprecate op related methods from AssetExecutionContext #19441
          • PR Make upstream asset materialization events available on the context #18971
            • PR [wip] store the tags fetched by data versioning so we can use them for upath io manager #19324
              • PR [UPath I/O managers] special case handling of None outputs #18820 👈
            • PR Make input partitions methods based on asset key not input name #19027
              • PR Support partitioned assets for fetching latest AssetMaterializations #19286
              • PR [do not merge] base example to showcase new apis #19199
                • PR [prototype] sub-context for each upstream dependency  #19283
                • PR [prototype] AssetExecutionContext partition methods on sub context #19236
                • PR [prototype] AssetExecutionContext partition methods given as named tuple #19037
                • PR [prototype] AssetExecutionContext partition methods directly on the context #19032
            • PR [wip] two execution contexts with minimal changes #18543
              • PR [wip] see wht needs to change when asset context is no longer an op context #18553
        • PR support Asset and Op ExecutionContexts in standard execution path #17972
          • PR [wip] add AssetExecutionContext to direct invocation #18044

This stack of pull requests is managed by Graphite.

[jamiedemaria]

Need to update has_output to check for the tag so that assets are supported for memoization, if supporting memoization for assets is even a thing we do

[jamiedemaria]

Putting the decision history for this PR's design here for future reference if we revive this effort:

Problem statement:

The UPath I/O manager stores None outputs in files and can load those None s as inputs in downstream assets. This causes two issues:

1. Potentially creates lots of None files on the FS, potentially confusing for users who aren’t using I/O managers
2. For users who aren’t using I/O managers, moving a non-partitioned asset to a partitioned asset will cause errors. A non-partitioned my_asset will make the file $DAGSTER_HOME/storage/my_asset but the partitioned asset will try to make the file $DAGSTER_HOME/storage/my_asset/2024-01-02 which will cause an error since my_asset is a file not a directory.

We want to modify the UPath I/O manager so that it would not create files for None outputs, but would still be able to load None s in downstream assets.

Solution 1: Use the existence of a file to determine if the asset was None . This has two problems:

1. The file not existing could mean the asset was never materialized. We’d need to query the event log to ensure the asset had been materialized before returning None
2. If an asset had been materialized with a non-None value, then later materialized with a None value, the original file would not be overwritten. At load time, the old non-None value would be loaded.
3. For ops, issues 1 and 2 don’t exist since the files are in a new directory for each run. But memoization relies on the existence of a file to determine if an op should rerun, so all ops that return None would always be rerun.

Solution 1 seemed bad, so we came up with…

Solution 2: Mark that the output was None on the corresponding AssetMaterialization / OutputHandled event

At load time, we can piggyback on the versioning code path to fetch the AssetMaterialization for the input assets. We can then check if this mark exists, and if so provide a None

Implementation 1: add this mark on the metadata for the AssetMaterialization /OutputHandled

Issues:

1. No code path exists for pre-fetching OutputHandled events for ops, so this would involve adding new “critical read” DB access during load_input and would require new indexes on the DB to get the OutputHandled event for a specific run id.
   1. solution - continue to store None returned from ops until memorization is deprecated and we can use file existence to determine None values.
2. in the versioning code path, the full AssetMaterialization is not fetched for partitioned assets, instead just the tags are fetched

Implementation 2: add this mark on the tags of the AssetMaterialization

This avoids issue 2 from above, but still has issues

1. The versioning code doesn’t fetch tags when more than 100 partitions are mapped for a particular input
   1. potential solution: When more than 100 partitions are mapped to an input, still fetch the none-output-tag in load_input . Would need to do this in batches so we don’t overwhelm the db. We can’t fetch the tag for each partition when we load the file for that particular partition since the load function is async, so we’d need to async-ify the db query, which we aren’t going to do.

More issues:

1. The versioning code is not called for asset checks, but we still need the None tag for asset checks. A potential solution is to also fetch the tags when an asset check is running (but not run the rest of the versioning code)
2. The defs.load_asset_value doesn’t have access to the tags either. We would need to fetch the tags when this is called.

---

If this does get revived, I'd recommend considering going directly into a deprecation cycle for storing None outputs. This metadata/tag shenanigans felt like it would add a lot of code smell and lead to confusion down the road as to what this code was for, and uncertainty about whether it could be deleted. It also relied heavily on assumptions about how AssetMaterializations are emitted for range-partition backfills. We require one AssetMaterialization per partition in the range, and that is a constraint that may not hold in the future. We can defend against this with unit tests, but it is still a fragility in the system and a non-intuitive code dependency.