How to parallelize assets testing in pytest ?

[louistransfer]

Hello,

I would like to have more information about the correct way to use multiprocessing in pytest for dagster assets.

I am currently trying to parallelize 3 unit tests which each assert the result of an asset polling an API. They have no dependencies to each other and can run in parallel. The processing time is quite long (roughly 5 minutes per test) as we have to wait for the API to process files, and it is the main bottleneck for CI completion.

Materialize to memory approach
Taking inspiration from this thread : #12346, I came up with the following code in my tests/test_api.py file :

def test_all_apis(load_env):

    api_client = APIClient()
    run_config = {
        "execution": {
            "config": {
                "multiprocess": {
                    "max_concurrent": 3, 
                }
            }
        }
    }
    with DagsterInstance.ephemeral() as instance:
        result = materialize_to_memory(
            [orders_api, sales_api, category_api],
            resources={"api_client": api_client},
            run_config=run_config,
            instance=instance,
        )
        assert result.success

However, the assets are still run sequentially.

Temporary reconstructable job approach
I tried another approach from this discussion #8467 by trying to create a reconstructable test job :

 def test_all_apis(load_env):
    api_client = APIClient()
    run_config = {
        "execution": {
            "config": {
                "multiprocess": {
                    "max_concurrent": 3,  # limits concurrent assets to 3
                }
            }
        }
    }

    test_apis_job = define_asset_job(name="test_apis_job", selection=[orders_api, sales_api, category_api])

    result = execute_job(
        reconstructable("test_apis_job"),
    instance=DagsterInstance.get(),
    run_config=run_config,
    )

    assert result.success

The issue here is that I can't see how I can declare this job in the DagsterInstance ? Would it be a good approach to run the 3 assets in parallel ?

Thanks for your help !

[alangenfeld]

Materialize to memory approach

If the assets are independent and I/O bound instead of compute bound, you can run the materialize_to_memory calls for each asset in threads for concurrency. Something like this:

def test_all_apis(load_env):
    api_client = APIClient()

    def _run_asset(a) -> bool:
        result = materialize_to_memory(
            [a],
            resources={"api_client": api_client},
        )
        return result.success

    with ThreadPoolExecutor() as executor:
        results = executor.map(_run_asset, [orders_api, sales_api, category_api])
        assert all(results)

Temporary reconstructable job approach

So the rules for reconstructable are just that it has to be able to find the job load function when we load up the step execution in another process. This can be accomplished by having a module scope function, something like this:

def get_test_job():
    assets_to_test = [orders_api, sales_api, category_api]
    test_job_name = "test_apis_job"
    defs = Definitions(
        jobs=[
            define_asset_job(
                name=test_job_name,
                selection=assets_to_test,
            )
        ],
        assets=assets_to_test,
        resources={"api_client": APIClient()},
    )
    return defs.get_job_def(test_job_name)


def test_all_apis():
    run_config = {
        "execution": {
            "config": {
                "multiprocess": {
                    "max_concurrent": 3,  # limits concurrent assets to 3
                }
            }
        }
    }

    with DagsterInstance.local_temp() as instance:
        result = execute_job(
            reconstructable(get_test_job),
            instance=instance,
            run_config=run_config,
        )

    assert result.success

[louistransfer]

Thanks for your detailed answer ! The first approach worked perfectly, our CI processing time was considerably reduced.

I think that your answer would fit perfectly in a section dedicated to multi-threading/parallel assets testing in the testing documentation, right after the section on graph assets :

https://docs.dagster.io/guides/dagster/testing-assets#testing-a-graph-of-assets

since I have seen other users expressing interest in this approach while searching for a solution.