Dagster Refactor: The Materialization Storm

We added a new dataset to our Dagster repo last week. It declared 50 dynamic-partition assets that depend on shared upstream assets. Dagster scheduler CPU pegged at 100% within an hour and we're now seeing 12,000 materialization runs per hour for what should be ~50. Run tests/test_no_runaway_runs.py to reproduce. Find the cycle/fanout pattern in dagster_repo/assets.py and fix it so the partition fan-out is bounded and the run count is sane. Be ready to defend the trade-off your fix makes.

dagster_repo/assets.py

# dagster_repo/assets.py
# Asset graph for the recommendations team. Recently added 50 per-tenant
# partitioned assets that hang off shared raw_events / user_profiles.
#
# Auto-materialize policy: any time an upstream materializes, every
# downstream asset that depends on it is queued.

from dagster_repo.kit import asset, AutoMaterializePolicy, DynamicPartitionsDefinition
from dagster_repo.io import read_events, read_profiles, write_recs

TENANTS = DynamicPartitionsDefinition(name='tenants')


@asset(auto_materialize_policy=AutoMaterializePolicy.eager())
def raw_events():
    return read_events()


@asset(auto_materialize_policy=AutoMaterializePolicy.eager())
def user_profiles():
    return read_profiles()


@asset(
    partitions_def=TENANTS,
    auto_materialize_policy=AutoMaterializePolicy.eager(),
    deps=[raw_events, user_profiles],
)
def tenant_events(context):
    tenant = context.partition_key
    events = raw_events()
    return [e for e in events if e['tenant'] == tenant]


@asset(
    partitions_def=TENANTS,
    auto_materialize_policy=AutoMaterializePolicy.eager(),
    deps=[tenant_events, user_profiles],
)
def tenant_features(context):
    return {'tenant': context.partition_key, 'feature_count': 5}


@asset(
    partitions_def=TENANTS,
    auto_materialize_policy=AutoMaterializePolicy.eager(),
    deps=[tenant_features],
)
def tenant_recommendations(context):
    return write_recs(context.partition_key)

# dagster_repo/assets.py
# Asset graph for the recommendations team. Recently added 50 per-tenant
# partitioned assets that hang off shared raw_events / user_profiles.
#
# Auto-materialize policy: any time an upstream materializes, every
# downstream asset that depends on it is queued.

from dagster_repo.kit import asset, AutoMaterializePolicy, DynamicPartitionsDefinition
from dagster_repo.io import read_events, read_profiles, write_recs

TENANTS = DynamicPartitionsDefinition(name='tenants')


@asset(auto_materialize_policy=AutoMaterializePolicy.eager())
def raw_events():
    return read_events()


@asset(auto_materialize_policy=AutoMaterializePolicy.eager())
def user_profiles():
    return read_profiles()


@asset(
    partitions_def=TENANTS,
    auto_materialize_policy=AutoMaterializePolicy.eager(),
    deps=[raw_events, user_profiles],
)
def tenant_events(context):
    tenant = context.partition_key
    events = raw_events()
    return [e for e in events if e['tenant'] == tenant]


@asset(
    partitions_def=TENANTS,
    auto_materialize_policy=AutoMaterializePolicy.eager(),
    deps=[tenant_events, user_profiles],
)
def tenant_features(context):
    return {'tenant': context.partition_key, 'feature_count': 5}


@asset(
    partitions_def=TENANTS,
    auto_materialize_policy=AutoMaterializePolicy.eager(),
    deps=[tenant_features],
)
def tenant_recommendations(context):
    return write_recs(context.partition_key)

Active Now|Staff Data Engineer (L6)|||3.6k Attempts|1.5k Solves|