How do you specify a custom storage path for each job run?

[eliu-fracta]

I have a job and a custom IO Manager defined below. What I want to achieve is the ability to specify a tag name when running the job, so that this tag name can later be used and read by the IO Manager as a keypath. For example, when I start a job run, I can specify a tag named "test_tag" and the IO Manager will add this tag as a prefix to all its storage paths, resulting in the format store/<test_tag>/pipes/cleaned_pipe_data.

IO manager

class LocalFileSystemIOManager(IOManager):

    def _get_fs_path(self, asset_key: AssetKey) -> str:
        rpath = os.path.join("/home/edward/Downloads/", *asset_key.path) + ".parquet"
        return os.path.abspath(rpath)

    def handle_output(self, context, obj: GeoDataFrame):
        fpath = self._get_fs_path(context.asset_key)
        context.log.info(fpath)
        context.add_output_metadata({"file path ": fpath})
        obj.to_parquet(fpath)

    def load_input(self, context):
        fpath = self._get_fs_path(context.asset_key)
        df = pd.read_parquet(fpath)
        if "geometry" in df.columns:
            df = gpd.read_parquet(fpath)

        return df

@io_manager
def parquet_io_manager():
    return LocalFileSystemIOManager()

Assets:

@asset(key_prefix=["store", "pipes"], compute_kind="pandas", io_manager_key="file_io")
def cleaned_pipe_data(context, config: CleanedPipeDataConfig):
    return df
    
@asset(key_prefix=["store", "breaks"], compute_kind="pandas", io_manager_key="file_io")
def cleaned_break_data(context, config: CleanedBreakDataConfig):
    return df


@asset(key_prefix=["store", "imputed"], compute_kind="pandas", io_manager_key="file_io")
def imputed_data(context, cleaned_pipe_data: gpd.GeoDataFrame, cleaned_break_data: pd.DataFrame):
    return cleaned_pipe_data    
    

Job:

etl_job = define_asset_job(
    name="etl_job", selection=["store/pipes/cleaned_pipe_data", "store/breaks/cleaned_break_data",
                               "store/imputed/imputed_data"],
    config=RunConfig(
        ops={
            "store__pipes__cleaned_pipe_data": CleanedPipeDataConfig(file_path="~/Downloads/Pipe_unimputed.parquet"),
            "store__breaks__cleaned_break_data": CleanedBreakDataConfig(file_path="~/Downloads/raw_komatsu_break.parquet")
        }
    ).to_config_dict())