[Python] Dagster: preserve multiple sklearn model in asset

I have a linear regression asset in dagster that uses data previously computed and sklearn LinearRegression (Python 10 here).

For each of my input columns (that represents a country, I want to fit a Linear Regression model.

Everything works fine. My question is about outputing these models (or use dagster metadata?)

I want basically a train asset and a forecast asset, a for this I want to return the models trained in the train asset and load them in the forecast asset. Solution could be to save them locally but I want to use dagster exclusively.

Also, I would like to save plenty of metadata (score, rmse) of each model into the train asset metadata.

Here is my code:

@asset(deps=[])
def train_linear_regression(duckdb: DuckDBResource):
"""Use pivot table with time serie data to forecast.
Used Linear Regression.
"""
# Setting up query.
query = "SELECT * FROM pivot_table_model"
# Execute the query.
with duckdb.get_connection() as conn:
df = conn.execute(query).df()
output = {}
for country_name in df.drop(columns=["year"]).columns:
# Setting Y.
Y = df.loc[:, country_name] # Retrieving population - pd.Series.
# Preparing linear model.
linear_regression = LinearRegression()
# Fitting the model.
linear_regression.fit(X, Y)
# Scoring.
score = linear_regression.score(X, Y)
output[country_name] = {
"model": linear_regression,
"score": float(score),
"plot": generate_plot(df),
}



What should I return? Any help appreciated

Regards,

Continue reading...