Distributed refactor

[LouisYRYJ]

This PR rewrites distributed setup to be more flexible.

[luciaquirke]

seems unnecessary

[luciaquirke]

if you can't see the keys for the input args to functions in your IDE something is broken, they should show up in grey

[luciaquirke]

I fixed this in the other query PR, I think it's a bit different than this fix. Basically returncode is 0 as long as the subprocess itself is happy so you parse the output text unconditionally

[luciaquirke]

set at init

[luciaquirke]

kwargs are undesirable when not totally necessary because they add a layer of indirection. if the SWE we're working with wants this I'm open to other lines of reasoning

[luciaquirke]

Need to add back the lost subset arg

[luciaquirke]

By the way, LM eval harness uses a nice pattern where the CLI takes --model_kwargs "device='cuda',streaming=True,subset='hello'" --dataset_kwargs ". . . " such that they don't need to update their library whenever HF adds a new model or dataset kwarg. We could consider doing this too.

[luciaquirke]

Need to add back the lost revision arg

[luciaquirke]

I'm removing this unless we have a good reason, I don't think we should suppress errors (?)

[luciaquirke]

I did this for the .part rename call and it was a mistake, I'm gonna remove mine

[luciaquirke]

Currently distributed_computing is not generic enough to be used in more than one place (it has build_index specific logic, like not having a QueryConfig parameter), and worker_fn is only ever set to collect_gradients, so it doesn't currently make sense to add a layer of indirection through an abstracted name there either.

worker_wrapper also has a generic name but does something specific (device-specific artifact setup & orchestration). I think the use of datasets is omnipresent in our workflows and we should embrace that as something to keep concrete. So we can name worker_wrapper to something like run_dataset_on_worker, and then we're not as surprised to find it's doing dataset processing.

Because both of our functions get the dataset in the main process, we can call maybe remove this conditional for now and add it back later if we need:

    # Do all the data loading and preprocessing on the main process
    if setup_data:
        ds = setup_data_pipeline(cfg)
    else:
        # Create empty dataset for compatibility
        ds = assert_type(Dataset, Dataset.from_list([]))

Then we may also choose to reduce the conceptual nesting of our functions by having a clear function called build that does

ds = setup_data_pipeline(cfg)
distributed_computing(worker_fn=collect_gradients, constant_worker_args=[cfg, ds], process_name="build"

and a query function that does

ds = setup_data_pipeline(index_cfg)
distributed_computing(worker_fn=collect_gradients, constant_worker_args=[index_cfg, query_cfg, ds], process_name="query"

Then distributed_computing will be truly generic, simply running

args={
i: (i, world_size, *constant_worker_args)
for i in range(world_size)
}

If we ever need to use something other than collect_gradients it can become one of the constant worker args.

I'm going to draft this up

[luciaquirke]

I'm going to yolo merge this, happy to raise another PR if there are issues or we want to add the kwargs pattern back to the tests or something like that

Thanks for working on this!!