Metadata not set (bug or user error?) on partitioned asset AssetMaterialization event?

[dboeckenhoff]

Hi everyone,New to dagster and I came up (together with your ai scout) with this mwe (which is still missing a lot of features).
What I try to accomplish is to trigger a graph job execution whenever a partition of my asset is materialized.
I use a AssetMaterialization event and its metadata for that (proposed by the ask-ai).

# content of w7x_dag/__init__.py
from typing import List
import subprocess
from dagster import (
    AssetKey,
    AssetMaterialization,
    ConfigurableResource,
    DynamicPartitionsDefinition,
    EventLogEntry,
    asset,
    asset_sensor,
    op,
    graph,
    In,
    Nothing,
    RunRequest,
    SensorEvaluationContext,
    Definitions,
)


# Define the API Resource
class ProgramAPI(ConfigurableResource):
    time_from: str = "2021-04-29 07:30:00"

    def fetch_programs(self, context=None) -> List[dict]:
        """
        This will return a list of program metadata (metadata on our experiments).
        """
        # This is a mock, for all times it will return ~5000 programs
        if self.time_from != "2021-04-29 07:30:00":
            raise NotImplementedError()

        return [
            {
                "description": "sniffer test - no gas, only short ECRH pulse\nQXT, QMR, QOI modified GeRI",
                "from": 1680158930574000000,
                "from_str": "1680158930574000000",
                "id": "20230330.001",
                "name": "#313707 sniffer_test_modGerI_C1B5",
                "upto": 1680159001238999999,
                "upto_str": "1680159001238999999"
            },
            {
                "description": "SOIII-67: ID10, ICRH Test",
                "from": 1680159427057000000,
                "from_str": "1680159427057000000",
                "id": "20230330.002",
                "name": "#404349 SOIII-67_ID10",
                "upto": 1680159499056999999,
                "upto_str": "1680159499056999999"
            }
        ]
    

partitions_def = DynamicPartitionsDefinition(name="program_partitions")


# Define Assets
@asset(partitions_def=partitions_def)
def programs(context, program_api: ProgramAPI):
    partition_key = context.partition_key
    programs = program_api.fetch_programs()
    context.log.info(f"Partion key: {partition_key}")
    for program in programs:
        if program["id"] == partition_key:
            context.log.info(f"Materializing program: {program}")
            context.log_event(
                AssetMaterialization(
                    asset_key="programs",
                    metadata={"test_string": "test", "program": program},
                )
            )
            return program
    context.log.error("No matching program found")
    return None


# Define Ops

"""
The ops (Analyses on our data) are standalone i.e. they don't require any
inputs but the metadata of an experiment (program) and write to our database.
"""


@op
def mock_analysis_low_level(context, program):
    subprocess.run(f"sleep 0.01 && echo 'Low ana {program['id']}'", shell=True)


@op
def mock_analysis_gather_artifacts(context, program):
    subprocess.run(f"sleep 0.011 && echo 'Artifacts {program['id']}'", shell=True)


@op(
    ins={
        "mock_analysis_low_level": In(Nothing),
        "mock_analysis_gather_artifacts": In(Nothing),
    }
)
def mock_analysis_high_level(context, program):
    subprocess.run(f"sleep 0.02 && echo 'High {program['id']}'", shell=True)


@graph
def process_programs_graph():
    low_level = mock_analysis_low_level()
    gather_artifacts = mock_analysis_gather_artifacts()
    mock_analysis_high_level(
        mock_analysis_low_level=low_level,
        mock_analysis_gather_artifacts=gather_artifacts,
    )


# Define jobs
process_programs_job = process_programs_graph.to_job()


# Define a Sensor
@asset_sensor(asset_key=AssetKey("programs"), job=process_programs_job)
def program_asset_sensor(context: SensorEvaluationContext, asset_event: EventLogEntry):
    program = asset_event.dagster_event.materialization.metadata["program"]
    context.log.info(f"Processing program: {program}")
    yield RunRequest(
        run_key=str(program["id"]),
        run_config={
            "ops": {
                "mock_analysis_low_level": {
                    "inputs": {"program": {"value": program}},
                },
                "mock_analysis_gather_artifacts": {
                    "inputs": {"program": {"value": program}},
                },
                "mock_analysis_high_level": {
                    "inputs": {"program": {"value": program}},
                },
            }
        },
    )


# Putting It All Together
defs = Definitions(
    assets=[programs],
    jobs=[process_programs_job],
    sensors=[program_asset_sensor],
    resources={"program_api": ProgramAPI()},
)


# Tests:
from dagster import build_sensor_context, materialize_to_memory
from dagster._core.test_utils import instance_for_test


def test_asset_sensor():
    program_api = ProgramAPI()
    with instance_for_test() as instance:
        # Add partitions to the dynamic partition definition
        instance.add_dynamic_partitions(partitions_def.name, ["20230330.001", "20230330.002"])
        # Materialize the programs asset
        materialize_to_memory(
            [programs],
            partition_key="20230330.001",
            resources={"program_api": program_api},
            instance=instance,
        )
        materialize_to_memory(
            [programs],
            partition_key="20230330.002",
            resources={"program_api": program_api},
            instance=instance,
        )
        context = build_sensor_context(
            instance=instance, resources={"program_api": program_api}
        )
        sensor_results = list(program_asset_sensor(context))

        assert len(sensor_results) == 2
        assert sensor_results[0].run_key == "20230330.001"
        assert sensor_results[1].run_key == "20230330.002"
        # Add assertions to verify the expected behavior

The test_asset_sensor does not pass because the metadata of the AssetMaterialization event seem to be dropped before they are passed to the program_asset_sensor. The metadata dict is empty.

❯ pytest "w7x_dag/__init__.py::test_asset_sensor"
=============================================================================================== test session starts ===============================================================================================
platform linux -- Python 3.10.14, pytest-8.3.2, pluggy-1.5.0
rootdir: /mnt/data/git/w7x_dag
configfile: pyproject.toml
plugins: time-machine-2.15.0, anyio-4.4.0, shutil-1.7.0, cov-5.0.0, virtualenv-1.7.0, xdist-3.6.1, hydra-core-1.3.2
collected 1 item                                                                                                                                                                                                  

w7x_dag/__init__.py F                                                                                                                                                                                       [100%]

==================================================================================================== FAILURES =====================================================================================================
________________________________________________________________________________________________ test_asset_sensor ________________________________________________________________________________________________

    def test_asset_sensor():
        program_api = ProgramAPI()
        with instance_for_test() as instance:
            # Add partitions to the dynamic partition definition
            instance.add_dynamic_partitions(partitions_def.name, ["20230330.001", "20230330.002"])
            # Materialize the programs asset
            materialize_to_memory(
                [programs],
                partition_key="20230330.001",
                resources={"program_api": program_api},
                instance=instance,
            )
            materialize_to_memory(
                [programs],
                partition_key="20230330.002",
                resources={"program_api": program_api},
                instance=instance,
            )
            context = build_sensor_context(
                instance=instance, resources={"program_api": program_api}
            )
>           sensor_results = list(program_asset_sensor(context))

w7x_dag/__init__.py:178: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
env/lib/python3.10/site-packages/dagster/_core/definitions/asset_sensor_definition.py:147: in _fn
    for item in result:
env/lib/python3.10/site-packages/dagster/_core/definitions/decorators/sensor_decorator.py:174: in _wrapped_fn
    for item in result:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

context = <dagster._core.definitions.sensor_definition.SensorEvaluationContext object at 0x7f5f75a278e0>
asset_event = EventLogEntry(error_info=None, level=10, user_message='Materialized value programs.', run_id='262d6205-391b-4cc1-8bb6-...0c155bd92b78576099a3a8'}), asset_lineage=[]), message='Materialized value programs.', pid=144455, step_key='programs'))

    @asset_sensor(asset_key=AssetKey("programs"), job=process_programs_job)
    def program_asset_sensor(context: SensorEvaluationContext, asset_event: EventLogEntry):
>       program = asset_event.dagster_event.materialization.metadata["program"]
E       KeyError: 'program'

w7x_dag/__init__.py:123: KeyError
---------------------------------------------------------------------------------------------- Captured stderr call -----------------------------------------------------------------------------------------------
2024-08-09 10:54:20 +0200 - dagster - DEBUG - __ephemeral_asset_job__ - fd8e7673-7dc4-4c45-a4d1-bab0c76908a7 - 144455 - RUN_START - Started execution of run for "__ephemeral_asset_job__".
2024-08-09 10:54:20 +0200 - dagster - DEBUG - __ephemeral_asset_job__ - fd8e7673-7dc4-4c45-a4d1-bab0c76908a7 - 144455 - ENGINE_EVENT - Executing steps in process (pid: 144455)
2024-08-09 10:54:20 +0200 - dagster - DEBUG - __ephemeral_asset_job__ - fd8e7673-7dc4-4c45-a4d1-bab0c76908a7 - 144455 - programs - RESOURCE_INIT_STARTED - Starting initialization of resources [io_manager, program_api].
2024-08-09 10:54:20 +0200 - dagster - DEBUG - __ephemeral_asset_job__ - fd8e7673-7dc4-4c45-a4d1-bab0c76908a7 - 144455 - programs - RESOURCE_INIT_SUCCESS - Finished initialization of resources [io_manager, program_api].
2024-08-09 10:54:20 +0200 - dagster - DEBUG - __ephemeral_asset_job__ - fd8e7673-7dc4-4c45-a4d1-bab0c76908a7 - 144455 - LOGS_CAPTURED - Started capturing logs in process (pid: 144455).
2024-08-09 10:54:20 +0200 - dagster - DEBUG - __ephemeral_asset_job__ - fd8e7673-7dc4-4c45-a4d1-bab0c76908a7 - 144455 - programs - STEP_START - Started execution of step "programs".
2024-08-09 10:54:20 +0200 - dagster - INFO - __ephemeral_asset_job__ - fd8e7673-7dc4-4c45-a4d1-bab0c76908a7 - programs - Partion key: 20230330.001
2024-08-09 10:54:20 +0200 - dagster - INFO - __ephemeral_asset_job__ - fd8e7673-7dc4-4c45-a4d1-bab0c76908a7 - programs - Materializing program: {'description': 'sniffer test - no gas, only short ECRH pulse\nQXT, QMR, QOI modified GeRI', 'from': 1680158930574000000, 'from_str': '1680158930574000000', 'id': '20230330.001', 'name': '#313707 sniffer_test_modGerI_C1B5', 'upto': 1680159001238999999, 'upto_str': '1680159001238999999'}
2024-08-09 10:54:20 +0200 - dagster - DEBUG - __ephemeral_asset_job__ - fd8e7673-7dc4-4c45-a4d1-bab0c76908a7 - 144455 - programs - ASSET_MATERIALIZATION - Materialized value programs.
2024-08-09 10:54:20 +0200 - dagster - DEBUG - __ephemeral_asset_job__ - fd8e7673-7dc4-4c45-a4d1-bab0c76908a7 - 144455 - programs - STEP_OUTPUT - Yielded output "result" of type "Any". (Type check passed).
2024-08-09 10:54:20 +0200 - dagster - DEBUG - __ephemeral_asset_job__ - fd8e7673-7dc4-4c45-a4d1-bab0c76908a7 - 144455 - programs - ASSET_MATERIALIZATION - Materialized value programs.
2024-08-09 10:54:20 +0200 - dagster - DEBUG - __ephemeral_asset_job__ - fd8e7673-7dc4-4c45-a4d1-bab0c76908a7 - 144455 - programs - HANDLED_OUTPUT - Handled output "result" using IO manager "io_manager"
2024-08-09 10:54:20 +0200 - dagster - DEBUG - __ephemeral_asset_job__ - fd8e7673-7dc4-4c45-a4d1-bab0c76908a7 - 144455 - programs - STEP_SUCCESS - Finished execution of step "programs" in 131ms.
2024-08-09 10:54:20 +0200 - dagster - DEBUG - __ephemeral_asset_job__ - fd8e7673-7dc4-4c45-a4d1-bab0c76908a7 - 144455 - ENGINE_EVENT - Finished steps in process (pid: 144455) in 260ms
2024-08-09 10:54:20 +0200 - dagster - DEBUG - __ephemeral_asset_job__ - fd8e7673-7dc4-4c45-a4d1-bab0c76908a7 - 144455 - RUN_SUCCESS - Finished execution of run for "__ephemeral_asset_job__".
2024-08-09 10:54:21 +0200 - dagster - DEBUG - __ephemeral_asset_job__ - 262d6205-391b-4cc1-8bb6-e525653c45cb - 144455 - RUN_START - Started execution of run for "__ephemeral_asset_job__".
2024-08-09 10:54:21 +0200 - dagster - DEBUG - __ephemeral_asset_job__ - 262d6205-391b-4cc1-8bb6-e525653c45cb - 144455 - ENGINE_EVENT - Executing steps in process (pid: 144455)
2024-08-09 10:54:21 +0200 - dagster - DEBUG - __ephemeral_asset_job__ - 262d6205-391b-4cc1-8bb6-e525653c45cb - 144455 - programs - RESOURCE_INIT_STARTED - Starting initialization of resources [io_manager, program_api].
2024-08-09 10:54:21 +0200 - dagster - DEBUG - __ephemeral_asset_job__ - 262d6205-391b-4cc1-8bb6-e525653c45cb - 144455 - programs - RESOURCE_INIT_SUCCESS - Finished initialization of resources [io_manager, program_api].
2024-08-09 10:54:21 +0200 - dagster - DEBUG - __ephemeral_asset_job__ - 262d6205-391b-4cc1-8bb6-e525653c45cb - 144455 - LOGS_CAPTURED - Started capturing logs in process (pid: 144455).
2024-08-09 10:54:21 +0200 - dagster - DEBUG - __ephemeral_asset_job__ - 262d6205-391b-4cc1-8bb6-e525653c45cb - 144455 - programs - STEP_START - Started execution of step "programs".
2024-08-09 10:54:21 +0200 - dagster - INFO - __ephemeral_asset_job__ - 262d6205-391b-4cc1-8bb6-e525653c45cb - programs - Partion key: 20230330.002
2024-08-09 10:54:21 +0200 - dagster - INFO - __ephemeral_asset_job__ - 262d6205-391b-4cc1-8bb6-e525653c45cb - programs - Materializing program: {'description': 'SOIII-67: ID10, ICRH Test', 'from': 1680159427057000000, 'from_str': '1680159427057000000', 'id': '20230330.002', 'name': '#404349 SOIII-67_ID10', 'upto': 1680159499056999999, 'upto_str': '1680159499056999999'}
2024-08-09 10:54:21 +0200 - dagster - DEBUG - __ephemeral_asset_job__ - 262d6205-391b-4cc1-8bb6-e525653c45cb - 144455 - programs - ASSET_MATERIALIZATION - Materialized value programs.
2024-08-09 10:54:21 +0200 - dagster - DEBUG - __ephemeral_asset_job__ - 262d6205-391b-4cc1-8bb6-e525653c45cb - 144455 - programs - STEP_OUTPUT - Yielded output "result" of type "Any". (Type check passed).
2024-08-09 10:54:21 +0200 - dagster - DEBUG - __ephemeral_asset_job__ - 262d6205-391b-4cc1-8bb6-e525653c45cb - 144455 - programs - ASSET_MATERIALIZATION - Materialized value programs.
2024-08-09 10:54:21 +0200 - dagster - DEBUG - __ephemeral_asset_job__ - 262d6205-391b-4cc1-8bb6-e525653c45cb - 144455 - programs - HANDLED_OUTPUT - Handled output "result" using IO manager "io_manager"
2024-08-09 10:54:21 +0200 - dagster - DEBUG - __ephemeral_asset_job__ - 262d6205-391b-4cc1-8bb6-e525653c45cb - 144455 - programs - STEP_SUCCESS - Finished execution of step "programs" in 134ms.
2024-08-09 10:54:21 +0200 - dagster - DEBUG - __ephemeral_asset_job__ - 262d6205-391b-4cc1-8bb6-e525653c45cb - 144455 - ENGINE_EVENT - Finished steps in process (pid: 144455) in 234ms
2024-08-09 10:54:21 +0200 - dagster - DEBUG - __ephemeral_asset_job__ - 262d6205-391b-4cc1-8bb6-e525653c45cb - 144455 - RUN_SUCCESS - Finished execution of run for "__ephemeral_asset_job__".
Test raised an exception, attempting to clean up instance:KeyError: 'program'

Stack Trace:
  File "/mnt/data/git/w7x_dag/env/lib/python3.10/site-packages/dagster/_core/instance_for_test.py", line 71, in instance_for_test
    yield instance
  File "/mnt/data/git/w7x_dag/w7x_dag/__init__.py", line 178, in test_asset_sensor
    sensor_results = list(program_asset_sensor(context))
  File "/mnt/data/git/w7x_dag/env/lib/python3.10/site-packages/dagster/_core/definitions/asset_sensor_definition.py", line 147, in _fn
    for item in result:
  File "/mnt/data/git/w7x_dag/env/lib/python3.10/site-packages/dagster/_core/definitions/decorators/sensor_decorator.py", line 174, in _wrapped_fn
    for item in result:
  File "/mnt/data/git/w7x_dag/w7x_dag/__init__.py", line 123, in program_asset_sensor
    program = asset_event.dagster_event.materialization.metadata["program"]

============================================================================================= short test summary info =============================================================================================
FAILED w7x_dag/__init__.py::test_asset_sensor - KeyError: 'program'
================================================================================================ 1 failed in 3.66s ================================================================================================

I am happy for all kinds of comments and advice (also not related to the question above).
Cheers!

[jamiedemaria]

hey @dboeckenhoff manually logging an asset materialization event should only be done in ops. All assets automatically log an asset materialization event when they complete, so what's happening is that that automatic materialization event is overriding the one you logged, hence the loss of metadata. Try logging just the metadata instead

@asset(partitions_def=partitions_def)
def programs(context, program_api: ProgramAPI):
    partition_key = context.partition_key
    programs = program_api.fetch_programs()
    context.log.info(f"Partion key: {partition_key}")
    for program in programs:
        if program["id"] == partition_key:
            context.log.info(f"Materializing program: {program}")
            context.add_output_metadata(
                {"test_string": "test", "program": program},
            )
            return program
    context.log.error("No matching program found")
    return None

This should work for you