I have an offline model pipeline that works 100% in spark: preprocessing, training LightGBM (synapse ml), and predictions. I am curious if this could fit into the UnionML framework. This would allow all the jobs to scale really nicely. I can see moving the preprocessing into the Dataset and then swapping out the DataFrame for a spark dataframe. There would need to be some config to have the tasks create spark jobs. Anyways let me know if this doesn’t sound crazy haha

This is not crazy at all! Need to improve the docs around this, but you can pass in all the

@task

kwargs to the following decorators: • Dataset.reader • Model.trainer • Model.predictor Meaning you can specify

SparkConfig

or any other Flyte-compatible task type configuration that typically work with Flyte tasks.

🤯 🤯 🤯

You can vary the features. Dataset.reader The dataset reader should output your dataset, and you can parameterize this function however way you wish. Here’s a basic example

@dataset.reader
def reader(feature_set: str, row_filter: str) -> pd.DataFrame
    data = ...  # get data using a SQL query, or whatever
    if feature_set == "feature_set_1":
        selected_data = ...
    elif feature_set == "feature_set_2":
        selected_data = ...
    else:
        ...  # etc
    
    if row_filter == "something":
        filtered_data = ...
    elif: ...

    return selected_data

So helpful!

Actually, would you mind creating an issue to document this use case? https://github.com/unionai-oss/unionml/issues/new Would love to capture it and add a page about this in the docs

Yeah of course!

Come to think of it, the Dataset.loader decorator might be a better place to vary datasets based on per-run parameters, but currently this isn’t supported. Working on this issue to address this tho

I added an issue and yes I agree to the dataset loader being a great place.