Retrieve data version of partition from downstream asset

[AbhinavPraveen]

I have an observable source asset like:

import dagster as dg
import os
import git

codes = dg.DynamicPartitionsDefinition(name="codes")

@dg.observable_source_asset(partitions_def=codes)
def code_partitions(context):
  code_url = "https://github.com/gitpython-developers/QuickStartTutorialFiles.git"
  repo = git.Repo.clone_from(code_url, "/tmp/repo")
  branches = [branch.name for branch in repo.branches]
  context.instance.add_dynamic_partitions(codes.name, branches)
  out = {branch: dg.DataVersion(repo.commit(branch).hexsha) for branch in branches}
  return dg.DataVersionsByPartition(out)

In a downstream partioned asset, during materialization, I want to clone exactly the commit and branch observed in the source asset. Is there a nice way to do this? I have read in #14723 that it may be possible to use the event log but it isn't obvious to me whether that works in the partioned case. Also, it seems like with this approach, it is possible that during a delay (e.g. Kubernetes pod scheduling), if another observation is taking on the source asset, the latest dataversion in the event log may change and thus the commit used in the downstream asset and the dataversion recorded in Dagster's database when the materialization is triggered may become inconsistent.

For example:

• Dagster observes the repo with a hash of 68b329da9893e34099c7d8ad5cb9c940
• A materialization of the downstream asset depending on the repo is trigerred and dagster records it against a repo data version of 68b329da9893e34099c7d8ad5cb9c940
• The job is scheduled but has not begun executing yet
• Dagster observes the repo again with a hash of b026324c6904b2a9cb4b88d6d61c81d1
• The job executes, retrieves the latest event and clones the repo at the commit b026324c6904b2a9cb4b88d6d61c81d1
• The commit used by the job and that recorded by dagster are inconsistent

I am not familiar with Dagster internals so please feel free to correct me if the above is incorrect for whatever reason. Any thoughts on whether there is a nice way to implement the above or to retrieve the data version of partition from downstream asset would be appreciated.

[apoorva-01]

Yeah, the event-log route from #14723 is the right instinct, it just isn't partition-aware. get_input_asset_record hands you the latest observation for the whole key. For a partitioned source you want to filter down to the current partition, which fetch_observations does:

from dagster import asset, AssetRecordsFilter

@dg.asset(partitions_def=codes, deps=[code_partitions])
def downstream(context):
    pk = context.partition_key

    records = context.instance.fetch_observations(
        AssetRecordsFilter(asset_key=code_partitions.key, asset_partitions=[pk]),
        limit=1,  # defaults to newest-first
    ).records

    commit = records[0].event_log_entry.tags.get("dagster/data_version")
    # == the hexsha you put in DataVersion(...) for this branch
    repo = git.Repo.clone_from(code_url, f"/tmp/{pk}")
    repo.git.checkout(commit)

The DataVersion you return from the observable source lands on the observation event as the dagster/data_version tag, so once you've grabbed the record for the right partition it's just a tag read.

One gotcha: this pulls whatever was observed last for that partition, not necessarily the version that kicked off this exact run. If observation and materialization go in lockstep you're fine. If observations can race ahead, pass the specific commit into the run (config or a run tag) instead of re-reading the newest.

[AbhinavPraveen]

Thanks, partly because of #33988, I ended up switching to assets instead and S3PickleIOManager to pass commits around. There doesn't seem to be an advantage to Observable Source Assets when using the OSS version.