https://github.com/stitchfix/hamilton logo
Join Slack
Powered by
Hi folks, trying to wrap my head around parallelis...
# hamilton-help
b
Hi folks, trying to wrap my head around parallelising a subdag. I'll place a dummy example as a thread response to keep the channel clean - cheers 🙂
Copy code
# site_etl.py
import pandas as pd
from hamilton.function_modifiers import config

def filter(product_df: pd.DataFrame, site: str) -> pd.DataFrame:
    return product_df[product_df['Site'] == site]

@config.when(sort_table=True)
def sort_table__true(filter: pd.DataFrame) -> pd.DataFrame:
    sort_df = filter.sort_values(by=['Date', 'Batch'], ascending=[True, True])
    return sort_df

@config.when(sort_table=False)
def sort_table__false(filter: pd.DataFrame) -> pd.DataFrame:
    return filter

@config.when(remove_duplicates=True)
def remove_table_dups__true(sort_table: pd.DataFrame, site: str) -> dict[str, pd.DataFrame]:
    rm_dup_df = sort_table.drop_duplicates(['Batch'])
    return {site: rm_dup_df}

@config.when(remove_duplicates=False)
def remove_table_dups__false(sort_table: pd.DataFrame, site: str) -> dict[str, pd.DataFrame]:
    return {site: sort_table}
Copy code
# subdag_etl.py

import site_etl
import pandas as pd
from hamilton.function_modifiers import subdag, source

@subdag(
    site_etl,
    inputs={'product_df': source('product_df')}
)

def etl_pipeline(remove_table_dups: dict[str, pd.DataFrame]) -> dict[str, pd.DataFrame]:
    return remove_table_dups
Copy code
# main.py

import subdag_etl
import pandas as pd
from hamilton import driver

data = [
    ("2021-01-31", "LT100", "US"),
    ("2022-02-21", "LT100", "US"),
    ("2022-01-02", "MT300", "US"),
    ("2020-04-14", "MT300", "US"),
    ("2020-09-10", "PX099", "CA"),
    ("2021-10-14", "PX099", "CA"),
    ("2014-12-10", "PX099", "CA")
]

product_df = pd.DataFrame(data, columns=['Date', 'Batch', 'Site'])
product_df['Date'] = pd.to_datetime(product_df['Date'], format="%Y-%m-%d").dt.normalize()

config = [
    {
        'site': 'US',
        'sort_table': True,
        'remove_duplicates': True
    },
    {
        'site': 'CA',
        'sort_table': False,
        'remove_duplicates': False
    }
]

for row in config:
        
    dr = (
        driver.Builder()
        .with_modules(subdag_etl)
        .with_config(config=row)
        .build()
    )

    result = dr.execute(['etl_pipeline'], inputs={'product_df': product_df})

    print(result['etl_pipeline'])

{'US':         Date  Batch Site
3 2020-04-14  MT300   US
0 2021-01-31  LT100   US}
{'CA':         Date  Batch Site
4 2020-09-10  PX099   CA
5 2021-10-14  PX099   CA
6 2014-12-10  PX099   CA}
My main source of confusion is when and where to extract the dictionaries in the configuration file and pass them in parallel to the subdag. In the past I have run a set of functions in parallel similarly to the 'graceful running' example by 1) reading a file, 2) iterating over the file (
yield()
and 
Parallelizable
), 3) Run n functions in parallel and 4) 
Collect
If I adopted the same strategy for the above, I would be removing the 
subdag_etl.py
file and thus the 
@subdag()
decorator entirely. Is there an elegant way to go about this?
s
@Barry Digby you want to parallelize this right:
```data = [
("2021-01-31", "LT100", "US"),
("2022-02-21", "LT100", "US"),
("2022-01-02", "MT300", "US"),
("2020-04-14", "MT300", "US"),
("2020-09-10", "PX099", "CA"),
("2021-10-14", "PX099", "CA"),
("2014-12-10", "PX099", "CA")
]```
?
If so, 
Parallelizable
is basically a way to "book end" a DAG. So you can place
Copy code
@subdag(
    site_etl,
    inputs={'product_df': source('product_df')}
)

def etl_pipeline(remove_table_dups: pd.DataFrame) -> pd.DataFrame:
    return remove_table_dups
in the middle of parallelizable & collect, and you should get what you want?
Some ideas here - "An advanced combination - Parallel with Subdag" (&

video for itâ–¾

)
b
Hi Stefan, My bad the dummy example was not clearer - I'll edit it now. The edited version shows my first attempt at this workflow where in 
main.py
I iterate over the config parameters, passing them to 
.with_config()
and execute the subdag sequentially in a for-loop. Works fine but I would like to run this in parallel - somewhere along the line I need either 
.with_config()
or 
config={}
to accept a 
Parallelizable
version of the config parameters. I think this deviates from the examples linked - they seem to be running input sources in parallel. In my case, 
product_df
remains fixed.
Copy code
# main.py --
# remove the for loop, pass config as dr.execute input for Parallelizable

dr = (
    driver.Builder()
    .enable_dynamic_execution(allow_experimental_mode=True)
    .with_remote_executor(executors.MultiThreadingExecutor(max_tasks=4))
    .with_modules(subdag_etl)
    .build()
)

result = dr.execute(['etl_pipeline'], inputs={'product_df': product_df, 'config': config})
Copy code
# subdag_etl.py --

def parse_config(config: list) -> Parallelizable[dict]:
    for row in config:
        yield row


@subdag(
    site_etl,
    inputs={'product_df': source('product_df')},
    config={?} 
)


def etl_pipeline(remove_table_dups: dict[str, pd.DataFrame]) -> dict[str, pd.DataFrame]:
    return remove_table_dups

    
def collect_results(etl_pipeline: Collect[dict[str, pd.DataFrame]]) -> dict:
    return etl_pipeline
I'm unsure how to pass the 
Parallelizable[dict]
to the 
@subdag(config=...)
- hope this makes sense, thank you for taking the time to review
s
ah - I misunderstood
configuration is resolved at driver build time
your only option I think is to nest Hamilton within Hamilton
or simply think of each config as an independent driver instantiation you can run independently however you'd like to parallelize it.
b
Hamilton within Hamilton did the trick, thanks for the help Stefan!
3 Views
Open in Slack
PreviousNext
Powered by