How to use the UploadDirectory plugin?

[josephnowak]

Hi,
I'm trying to add one of the local modules of my API to the Dask workers with the UploadDirectory plugin but I'm getting a ModuleNotFoundError, could you indicate to me what I'm doing wrong? I have created the following small project on a Jupyter notebook to exemplify in a simple way what I'm doing on the API project.

This is the structure:

This is the code on the main file which is invocating the UploadDirectory:

from dask.distributed import Client
from distributed.diagnostics.plugin import UploadDirectory

client = Client("CLUSTER_URL")
plugin = UploadDirectory("/check_imports", restart=True, update_path=True)
client.register_plugin(plugin)
# I'm not sure if the restart method is necessary
client.restart()

def f():
    from check_imports.func import get_func_data    
    return get_func_data()

f().compute()

This is the code from the other_func.py file inside the check_imports folder:

from typing import Optional

import dask.dataframe as dd
import pandas as pd


def get_other_func_data(df) -> pd.DataFrame:
    return df.sum()


def get_delayed_df():
    df = dd.from_pandas(pd.DataFrame([4] * 100), chunksize=10)
    df = df.map_partitions(get_other_func_data)
    return df

Code from the func.py file inside the check_imports folder:

import dask.dataframe as dd
import pandas as pd
from .other_func import get_delayed_df


def get_func_data() -> pd.DataFrame:
    df = get_delayed_df()
    return df