Skip to content

[WORKAROUND] Traditional schedule for e assets, eager for h* assets #1060

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
damonbayer merged 7 commits into from
May 12, 2026
Merged

[WORKAROUND] Traditional schedule for e assets, eager for h* assets #1060

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
453ce60
disable upstream sensor; add schedule for e assets
jkislin May 11, 2026
13d36ee
weekly forecast H and weekly forecast E groups that we can supply new…
jkislin May 11, 2026
a0b484f
some tidying and bug fixes
jkislin May 11, 2026
2e82248
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] May 11, 2026
c6357bb
Merge branch 'main' into jk-schedule-workaround
jkislin May 11, 2026
b83f70b
Merge branch 'main' into jk-schedule-workaround
jkislin May 12, 2026
f94ce3e
submitting pyrenew-e jobs wednesday only
jkislin May 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 11 additions & 7 deletions EpiAutoGP/test/test_output.jl
View file
Open in desktop

@giomrella giomrella May 12, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This seems unintentional

@jkislin jkislin May 12, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's a pre-commit change.

@damonbayer damonbayer May 12, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

working on getting julia added to pre-commit.ci so this doesn't happen pre-commit-ci/runner-image#332

Original file line number Diff line number Diff line change
Expand Up @@ -114,23 +114,27 @@ end
Symbol(".variable"),
:resolution,
:geo_value,
:disease
:disease,
]

con = DBInterface.connect(DuckDB.DB, ":memory:")
try
read_df = DataFrame(DBInterface.execute(
con,
"SELECT * FROM read_parquet($(EpiAutoGP._quote_duckdb_string(parquet_path)))"
))
read_df = DataFrame(
DBInterface.execute(
con,
"SELECT * FROM read_parquet($(EpiAutoGP._quote_duckdb_string(parquet_path)))"
)
)

@test eltype(read_df.date) == Date
@test propertynames(read_df) == propertynames(result_df)
@test read_df.date == forecast_dates[[1, 2, 1, 2]]
@test read_df[!, Symbol(".draw")] == [1, 1, 2, 2]
@test read_df[!, Symbol(".value")] == [10.0, 20.0, 30.0, 40.0]
@test all(read_df[!, Symbol(".variable")] .==
"observed_hospital_admissions")
@test all(
read_df[!, Symbol(".variable")] .==
"observed_hospital_admissions"
)
@test all(read_df.resolution .== "epiweekly")
@test all(read_df.geo_value .== "CA")
@test all(read_df.disease .== "COVID-19")
Expand Down
159 changes: 102 additions & 57 deletions dagster_defs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -459,10 +459,55 @@ def _fuse_pyrenew_timeseries(
# SCHEDULES AND AUTOMATION CONDITION SENSORS
# ============================================================================

weekly_forecast_upstream_sensor = dg.AutomationConditionSensorDefinition(
name="WeeklyForecastUpstream",
target=dg.AssetSelection.groups("WeeklyForecastUpstream"),
use_user_code_server=True, # allows for custom automation conditions
# Custom Automation Condition. Relies on use_user_code_server=True on the sensor
# class IsWeekday(dg.AutomationCondition):
# def __init__(self, weekday: int):
# """
# Check if evaluation time falls on a specific weekday.
# This is is a simple evaluation, rather than a stateful operation,
# such as with cron_tick_passed().

# Args:
# weekday: 0=Monday, 1=Tuesday, 2=Wednesday, 3=Thursday,
# 4=Friday, 5=Saturday, 6=Sunday
# """
# self.weekday = weekday
# super().__init__()

# def evaluate(self, context: dg.AutomationContext) -> dg.AutomationResult:
# # If the current weekday is equal to the desired weekday,
# # return the candidate_subset -> a dagster context's "true" case
# if context.evaluation_time.weekday() == self.weekday:
# true_subset = context.candidate_subset
# else:
# true_subset = context.get_empty_subset()

# return dg.AutomationResult(context=context, true_subset=true_subset)

# @property
# def name(self) -> str:
# """Define the label that will appear in the UI"""
# days = [
# "Monday",
# "Tuesday",
# "Wednesday",
# "Thursday",
# "Friday",
# "Saturday",
# "Sunday",
# ]
# return f"is_{days[self.weekday].lower()}"

# weekly_forecast_upstream_sensor = dg.AutomationConditionSensorDefinition(
# name="WeeklyForecastUpstream",
# target=dg.AssetSelection.groups("WeeklyForecastUpstream"),
# use_user_code_server=True, # allows for custom automation conditions
# )

Comment thread
jkislin marked this conversation as resolved.
weekly_forecast_h_sensor = dg.AutomationConditionSensorDefinition(
name="WeeklyForecastH",
target=dg.AssetSelection.groups("WeeklyForecastH"),
use_user_code_server=False, # does NOT allow custom conditions
)

weekly_forecast_fusion_sensor = dg.AutomationConditionSensorDefinition(
Expand All @@ -472,44 +517,28 @@ def _fuse_pyrenew_timeseries(
)


# Custom Automation Condition. Relies on use_user_code_server=True on the sensor
class IsWeekday(dg.AutomationCondition):
def __init__(self, weekday: int):
"""
Check if evaluation time falls on a specific weekday.
This is is a simple evaluation, rather than a stateful operation,
such as with cron_tick_passed().

Args:
weekday: 0=Monday, 1=Tuesday, 2=Wednesday, 3=Thursday,
4=Friday, 5=Saturday, 6=Sunday
"""
self.weekday = weekday
super().__init__()

def evaluate(self, context: dg.AutomationContext) -> dg.AutomationResult:
# If the current weekday is equal to the desired weekday,
# return the candidate_subset -> a dagster context's "true" case
if context.evaluation_time.weekday() == self.weekday:
true_subset = context.candidate_subset
else:
true_subset = context.get_empty_subset()

return dg.AutomationResult(context=context, true_subset=true_subset)

@property
def name(self) -> str:
"""Define the label that will appear in the UI"""
days = [
"Monday",
"Tuesday",
"Wednesday",
"Thursday",
"Friday",
"Saturday",
"Sunday",
]
return f"is_{days[self.weekday].lower()}"
# Temporary workaround schedule while we debug the custom automation condition
@dg.schedule(
target=dg.AssetSelection.assets(
"timeseries_e", "epiweekly_timeseries_e", "pyrenew_e"
),
cron_schedule="30 6 * * WED", # 6:30am on Wednesday (day 3)
execution_timezone="America/New_York",
)
def weekly_forecast_e_schedule(context: dg.ScheduleEvaluationContext):
_partition_key = daily_partitions_def.get_last_partition_key()
context.log.info(f"Submitting job request for partition: {_partition_key}")
return dg.RunRequest(
partition_key=_partition_key,
run_config=dg.RunConfig(
Comment thread
jkislin marked this conversation as resolved.
ops={
"timeseries_e": TimeseriesConfig(),
"epiweekly_timeseries_e": TimeseriesConfig(),
"pyrenew_e": PyrenewEConfig(),
},
execution=azure_batch_execution_config.to_run_config(),
),
)


# ---------- Shared Asset Decorator Arguments ----------
Expand All @@ -518,22 +547,26 @@ def name(self) -> str:
# partitions, graph_dimensions, automation conditions, and asset groups
# The only thing that differs between them are their dependencies

weekly_forecast_upstream_asset_args = {
weekly_forecast_base_asset_args = {
"partitions_def": daily_partitions_def,
"graph_dimensions": ["diseases", "locations"],
"group_name": "WeeklyForecastUpstream",
"automation_condition": (
# We specifically don't want these to run unless it's Wednesday
# 0=monday,1=tuesday,2=wednesday,etc.
# Note this is different from cron which is 1-indexed
dg.AutomationCondition.eager() & IsWeekday(2)
).with_label("eager_on_wednesday"),
}

# weekly_forecast_upstream_asset_args = {
# **weekly_forecast_base_asset_args,
# "group_name": "WeeklyForecastUpstream",
# "automation_condition": (
# # We specifically don't want these to run unless it's Wednesday
# # 0=monday,1=tuesday,2=wednesday,etc.
# # Note this is different from cron which is 1-indexed
# dg.AutomationCondition.eager() & IsWeekday(2)
# ).with_label("eager_on_wednesday"),
# }

weekly_forecast_fusion_asset_args = {
**weekly_forecast_upstream_asset_args,
**weekly_forecast_base_asset_args,
"group_name": "WeeklyForecastFusion",
# we want vanilla eager for the fusion assets and post-processing
# we want vanilla eager for the fusion assets
"automation_condition": dg.AutomationCondition.eager(),
}

Expand Down Expand Up @@ -564,7 +597,9 @@ def name(self) -> str:

# Timeseries E
@dynamic_graph_asset(
**weekly_forecast_upstream_asset_args,
# **weekly_forecast_upstream_asset_args,
**weekly_forecast_base_asset_args,
group_name="WeeklyForecastE", # This will override what's in the asset args for now
ins={"nssp_gold_v1": dg.In(dg.Nothing)},
)
def timeseries_e(context: DynamicGraphAssetExecutionContext, config: TimeseriesConfig):
Expand All @@ -573,7 +608,9 @@ def timeseries_e(context: DynamicGraphAssetExecutionContext, config: TimeseriesC

# Epiweekly Timeseries E
@dynamic_graph_asset(
**weekly_forecast_upstream_asset_args,
# **weekly_forecast_upstream_asset_args,
**weekly_forecast_base_asset_args,
group_name="WeeklyForecastE", # This will override what's in the asset args for now
ins={"nssp_gold_v1": dg.In(dg.Nothing)},
)
def epiweekly_timeseries_e(
Expand All @@ -584,7 +621,9 @@ def epiweekly_timeseries_e(

# Pyrenew E
@dynamic_graph_asset(
**weekly_forecast_upstream_asset_args,
# **weekly_forecast_upstream_asset_args,
**weekly_forecast_base_asset_args,
group_name="WeeklyForecastE", # This will override what's in the asset args for now
ins={
"nssp_gold_v1": dg.In(dg.Nothing),
},
Expand All @@ -598,7 +637,10 @@ def pyrenew_e(

# Pyrenew H
@dynamic_graph_asset(
**weekly_forecast_upstream_asset_args,
# **weekly_forecast_upstream_asset_args,
**weekly_forecast_base_asset_args,
automation_condition=dg.AutomationCondition.eager(), # H assets can be eager, inheriting their schedule from dataops
group_name="WeeklyForecastH", # This will override what's in the asset args for now
ins={
"nhsn_hrd_prelim": dg.In(dg.Nothing),
},
Expand All @@ -609,7 +651,10 @@ def pyrenew_h(context: DynamicGraphAssetExecutionContext, config: PyrenewConfig)

# Pyrenew HE
@dynamic_graph_asset(
**weekly_forecast_upstream_asset_args,
# **weekly_forecast_upstream_asset_args,
**weekly_forecast_base_asset_args,
automation_condition=dg.AutomationCondition.eager(), # H assets can be eager, inheriting their schedule from dataops
group_name="WeeklyForecastH", # This will override what's in the asset args for now
ins={
"nhsn_hrd_prelim": dg.In(dg.Nothing),
},
Expand Down
Loading