[Feature] Add data validation steps to `trainer.dev` for early error detection

[dalek-who]

Description

Currently, data conversion steps (e.g., converting data to PyArrow arrays) only run inside trainer.fit. This causes two issues:

1. Data format errors are only caught during training, not during development.
2. When debugging in VSCode, breakpoints set in python libs (e.g., arrow_writer.py) are not hit during trainer.dev, making debugging harder.

Example

In trainer.fit, one conversion step is:

# From: /root/miniconda3/envs/agl/lib/python3.12/site-packages/datasets/arrow_writer.py
out = pa.array(cast_to_python_objects(data, only_1d_for_numpy=True))

If the data format is incorrect, it may raise an error, but the breakpoint here will not hit.

Proposed Solution

Add the same data validation/conversion steps of trainer.fit to trainer.dev. This will:

1. Catch data format errors earlier.
2. Allow breakpoints in data processing code to be triggered during development.

[ultmaster]

Could you elaborate on "converting data to PyArrow arrays"? Where did our code do such conversion?

[dalek-who]

Well, it appears that this is related to a specific algorithm (verl in this case), and cannot be directly ported to trainer.dev.

  File "/workspace/projects/agl-reflection/./train.py", line 88, in train
    trainer.fit(agent, train_dataset=train_data, val_dataset=val_data)  # type: ignore
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/miniconda3/envs/agl/lib/python3.12/site-packages/agentlightning/trainer/trainer.py", line 424, in fit
    self.strategy.execute(algorithm_bundle, runner_bundle, self.store)
  File "/root/miniconda3/envs/agl/lib/python3.12/site-packages/agentlightning/execution/client_server.py", line 386, in execute
    asyncio.run(self._execute_algorithm(algorithm, store, stop_evt))
  File "/root/miniconda3/envs/agl/lib/python3.12/asyncio/runners.py", line 195, in run
    return runner.run(main)
           ^^^^^^^^^^^^^^^^
  File "/root/miniconda3/envs/agl/lib/python3.12/asyncio/runners.py", line 118, in run
    return self._loop.run_until_complete(task)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/miniconda3/envs/agl/lib/python3.12/asyncio/base_events.py", line 691, in run_until_complete
    return future.result()
           ^^^^^^^^^^^^^^^
  File "/root/miniconda3/envs/agl/lib/python3.12/site-packages/agentlightning/execution/client_server.py", line 154, in _execute_algorithm
    await algorithm(wrapper_store, stop_evt)
  File "/root/miniconda3/envs/agl/lib/python3.12/site-packages/agentlightning/trainer/trainer.py", line 513, in _algorithm_bundle
    algorithm.run(
  File "/root/miniconda3/envs/agl/lib/python3.12/site-packages/agentlightning/algorithm/verl/interface.py", line 138, in run
    run_ppo(
  File "/root/miniconda3/envs/agl/lib/python3.12/site-packages/agentlightning/verl/entrypoint.py", line 58, in run_ppo
    ray.get(
  File "/root/miniconda3/envs/agl/lib/python3.12/site-packages/ray/_private/auto_init_hook.py", line 22, in auto_init_wrapper
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/root/miniconda3/envs/agl/lib/python3.12/site-packages/ray/_private/client_mode_hook.py", line 104, in wrapper
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/root/miniconda3/envs/agl/lib/python3.12/site-packages/ray/_private/worker.py", line 2972, in get
    values, debugger_breakpoint = worker.get_objects(
                                  ^^^^^^^^^^^^^^^^^^^
  File "/root/miniconda3/envs/agl/lib/python3.12/site-packages/ray/_private/worker.py", line 1031, in get_objects
    raise value.as_instanceof_cause()
ray.exceptions.RayTaskError(ArrowInvalid): ray::TaskRunner.run() (pid=34185, ip=172.17.0.2, actor_id=69ab9382ce02d10b49aed87801000000, repr=<agentlightning.verl.entrypoint.TaskRunner object at 0x7eff74392540>)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/miniconda3/envs/agl/lib/python3.12/site-packages/agentlightning/verl/entrypoint.py", line 190, in run
    val_dataset = LoadedDataset(val_dataset)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/miniconda3/envs/agl/lib/python3.12/site-packages/agentlightning/verl/dataset.py", line 41, in __init__
    self.dataframe = HuggingFaceDataset.from_list(dataset_copy)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/miniconda3/envs/agl/lib/python3.12/site-packages/datasets/arrow_dataset.py", line 1068, in from_list
    return cls.from_dict(mapping, features, info, split)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/miniconda3/envs/agl/lib/python3.12/site-packages/datasets/arrow_dataset.py", line 1022, in from_dict
    pa_table = InMemoryTable.from_pydict(mapping=mapping)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/miniconda3/envs/agl/lib/python3.12/site-packages/datasets/table.py", line 757, in from_pydict
    return cls(pa.Table.from_pydict(*args, **kwargs))
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "pyarrow/table.pxi", line 1985, in pyarrow.lib._Tabular.from_pydict
  File "pyarrow/table.pxi", line 6401, in pyarrow.lib._from_pydict
  File "pyarrow/array.pxi", line 405, in pyarrow.lib.asarray
  File "pyarrow/array.pxi", line 256, in pyarrow.lib.array
  File "pyarrow/array.pxi", line 118, in pyarrow.lib._handle_arrow_array_protocol
  File "/root/miniconda3/envs/agl/lib/python3.12/site-packages/datasets/arrow_writer.py", line 308, in __arrow_array__
    out = pa.array(cast_to_python_objects(data, only_1d_for_numpy=True))
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "pyarrow/array.pxi", line 375, in pyarrow.lib.array
  File "pyarrow/array.pxi", line 46, in pyarrow.lib._sequence_to_array
  File "pyarrow/error.pxi", line 155, in pyarrow.lib.pyarrow_internal_check_status
  File "pyarrow/error.pxi", line 92, in pyarrow.lib.check_status
pyarrow.lib.ArrowInvalid: cannot mix list and non-list, non-null values

[ultmaster]

but you are using trainer.fit? What's the relationship between that and trainer.dev?

Maybe you should convert the data into a native list before passing it to trainer?

[dalek-who]

The above error originates from trainer.fit, not trainer.dev. My data has already been converted to a native list. The error is due to some inherent issues in the data itself (unrelated to agentlightning).
My point is that the data processing steps (such as conversion to pyarrow) occur in trainer.fit, but not in trainer.dev. However, if this process could also be included in trainer.dev, data issues could be detected earlier before training starts (since when running trainer.fit, even setting a breakpoint at the error location /root/miniconda3/envs/agl/lib/python3.12/site-packages/datasets/arrow_writer.py won’t pause the execution).

[ultmaster]

Could you maybe share what kind of data you are dealing with so that I can reproduce the issue locally?

[dalek-who]

Here is the simplest example:

val_data_1 = [
  {"input": "1+1=", "target": ["two", "2"]},
  {"input": "1+2=", "target": ["three", "3"]},
  {"input": "3+2=", "target": ["five", "5"]},
]  

val_data_2 = [
  {"input": "1+1=", "target": ["two", "2"]},
  {"input": "1+2=", "target": ["three", "3"]},
  {"input": "3+2=", "target": "5"}
]

For trainer.fit, val_data_1 does not trigger an error, while val_data_2 triggers anpyarrow.lib.ArrowInvalid: cannot mix list and non-list, non-null values when using the verl algorithm. The reason is that the "target" in {"input": "3+2=", "target": "5"} is a string, while the other "target" fields are lists.

Although this example is artificial, similar situations can occur when merging data from different sources, where inconsistencies in format during preprocessing go unnoticed, and are only discovered during training.

[ultmaster]

I think it's because the data is first converted to some huggingface dataset within verl. It's a verl limitation but it's possible we break it, as we already rewrite a new dataset format (not verl.RLHFDataset any more).