[data][llm] Add per-stage map kwargs for build_llm_processor preprocess/postprocess

[nrghosh]

Description

Enable users to control resources and concurrency of preprocess/postprocess stages independently from the main LLM stage by passing Dataset.map() kwargs.

• Add preprocess_map_kwargs and postprocess_map_kwargs parameters to build_llm_processor() and all builder functions
• Update Processor class to store and apply map kwargs to dataset.map() calls
• Add validation for map kwargs with warnings on unknown keys
• Unit tests + Update docstrings

Users can now provision fractional CPU resources and tune parameters per stage without workarounds- improving utilization.

Related issues

Addresses #57812

Additional information

Example usage:

  processor = build_llm_processor(
      config,
      preprocess=preprocess_fn,
      postprocess=postprocess_fn,
      preprocess_map_kwargs={"num_cpus": 0.5},
      postprocess_map_kwargs={"num_cpus": 0.25},
  )

instead of

# preprocess 
ds = ray_dataset.map_batches(
    preprocess,  
    batch_size=1,  
    num_cpus=0.5,
    num_gpus=0,
)

# processor with preprocess=None, postprocess=None
proc = build_llm_processor(processor_config, preprocess=None, postprocess=None)

# vLLM stage
generated = proc(ds)

# Postprocess
captioned = generated.map_batches(
    postprocess,
    num_cpus=0.25,   
    num_gpus=0,
)

[kouroshHakha]

this list is dynamic and i don't think it's easy to maintain this one. We should just leave it up to the user to read the docs and provide the right map kwargs.

[nrghosh]

good point, resolved

[kouroshHakha]

@nrghosh tests are failing.