feat(api): (1/n) datasets api clean up

[yanxi0830]

PR Stack

• feat(api): (1/n) datasets api clean up #1573
• feat(dataset api): (1.1/n) dataset api implementation fix pre-commit #1625
• feat(dataset api): (1.2/n) datasets.iterrorws pagination api updates #1656
• feat(dataset api): (1.3/n) patch OpenAPI gen for datasetio->datasets #1657
• feat(dataset api): (1.4/n) fix resolver signature mismatch #1658
• feat(dataset api): (1.5/n) fix dataset registeration #1659
• feat(dataset api): (1.6/n) fix all iterrows callsites #1660

Client SDK

• feat: datasets api updates llama-stack-client-python#203

CI

• https://github.com/meta-llama/llama-stack-ops/actions/runs/13911304888

-- the test_rag_agent_with_attachments is flaky and not related to this PR

Doc

Client Usage

client.datasets.register(
    source={
        "type": "uri",
        "uri": "lsfs://mydata.jsonl",
    },
    schema="jsonl_messages",
    # optional 
    dataset_id="my_first_train_data"
)

# quick prototype debugging
client.datasets.register(
    data_reference={
        "type": "rows",
        "rows": [
                "messages": [...],
        ],
    },
    schema="jsonl_messages",
)

Test Plan

• CI: https://github.com/meta-llama/llama-stack-ops/actions/runs/13878055453

LLAMA_STACK_CONFIG=fireworks pytest -v tests/integration/datasets/test_datasets.py

LLAMA_STACK_CONFIG=fireworks pytest -v tests/integration/scoring/test_scoring.py

pytest -v -s --nbval-lax ./docs/notebooks/Llama_Stack_Benchmark_Evals.ipynb

[ashwinb]

when would file:// make sense? when the server is running on the same machine as client? i'd rather not have that.

[yanxi0830]

I believe the idea was to use file:// to indicate that the dataset will be coming from a File object uploaded via file API.

• file://<file-id> -> loads data from local file system via file ID
• huggingface://<dataset-path> -> loads data from hugginface

[yanxi0830]

The other alternative is instead of a single uri, we have different Data class to specify different ways we can create a dataset.

class DatasetType(Enum):
    huggingface = "huggingface"
    file = "file"

class HuggingfaceData(BaseModel):
    type: Literal[DatasetType.huggingface.value]
    dataset_id: str
    split: str

class FileData(BaseModel):
    type: Literal[DatasetType.file.value]
    file_id: str

Data = Annotated[
    Union[FileData, HuggingfaceData],
    Field(discriminator="type"),
]

 async def register_dataset(
        self,
        schema: Schema,
        data: Data,
        metadata: Optional[Dict[str, Any]] = None,
        dataset_id: Optional[str] = None,
    ) -> Dataset:

[yanxi0830]

Synced with @hardikjshah offline. Going with the Union way with data_reference data_source.

[hardikjshah]

maybe just source, drop the data_ ? This is inside of a dataset anyways

[yanxi0830]

update to source

[ashwinb]

I don't think we should have file:// as an allowed URI. if you want local files, you can calculate rows yourself and pass as rows. Otherwise you are relying on the client-SDK faithfully obeying this API (and we will not be able to keep this in sync on all clients without additional effort)

[yanxi0830]

I don't think we should have file:// as an allowed URI.

@ashwinb Added example URI's that is allowed

- "https://mywebsite.com/mydata.jsonl"
- "lsfs://mydata.jsonl" --> a file url
- "data:csv;base64,{base64_content}"

[ehhuang]

A bit confused here. Does this dictate what the format of the data source should be? e.g. can I supply a csv file of type uri as data source?

[yanxi0830]

The Schema defines the column name and format. For example:

1. messages schema indicate a dataset consisting of rows with a "messages" column with List[Messages]. These datasets can be used as input for finetuning.

[
    {"messages": List[Message]}
]

2. chat_completion_eval consisting of "messages" column and "expected_answer" column. The "messages" column will be used as input to generate "generated_answer", and "expected_answer" column will be used for eval for scoring and computing metrics.

The reason we need schema is to validate that the supplied dataset can be correctly parsed by different methods:

llama-stack/llama_stack/providers/utils/common/data_schema_validator.py

Lines 28 to 53 in 00da911

	VALID_SCHEMAS_FOR_SCORING = [
	{
	ColumnName.input_query.value: StringType(),
	ColumnName.expected_answer.value: StringType(),
	ColumnName.generated_answer.value: StringType(),
	},
	{
	ColumnName.input_query.value: StringType(),
	ColumnName.expected_answer.value: StringType(),
	ColumnName.generated_answer.value: StringType(),
	ColumnName.context.value: StringType(),
	},
	]
	VALID_SCHEMAS_FOR_EVAL = [
	{
	ColumnName.input_query.value: StringType(),
	ColumnName.expected_answer.value: StringType(),
	ColumnName.chat_completion_input.value: ChatCompletionInputType(),
	},
	{
	ColumnName.input_query.value: StringType(),
	ColumnName.expected_answer.value: StringType(),
	ColumnName.completion_input.value: CompletionInputType(),
	},
	]

The DataSource represent where the data can be loaded from. A csv file can be supplied via 3 ways:

1. base64 encoded in the uri.

URIDataSource(
     uri="`data:csv;base64,{base64_content}`",
)

2. upload csv file vie /files API defined in https://github.com/meta-llama/llama-stack/blob/main/llama_stack/apis/files/files.py

3. read rows in csv file and register dataset via:

client.datasets.create(
    data_reference={
        "type": "rows",
        "rows": [
                "messages": [...],
        ],
    },
    schema="jsonl_messages",
)

[yanxi0830]

Synced with @ehhuang offline. purpose might be a better name compared with schema and less confusing for users, since we could infer the schema from dataset. cc @hardikjshah for thoughts.

[ehhuang]

Discussed offline:

• Clarification: The purpose of the schema parameter is to validate that the input data has the expected format.
• Perhaps a more apt name would be purpose, where in the documentation we lay out the expected schema for each purpose, and in the implementation of create we validate that the input data against the expected schema.

[hardikjshah]

What are all the valid schema names that we will have given what we know right now ?

• jsonl with each line being a `{"messages": []} ( tuning/messages )
• jsonl with each line being {"text": ... } ( tuning/text )
• eval with a csv ( query / expected / output ) ( eval/question_answer )

May be we use this api/purpose style to name the various schemas ?

[ashwinb]

nit: this is not flaky (which means it sometimes works vs. sometimes not) -- I think maybe you just mean "wonkiness" or "sadness"

[yanxi0830]

yeah, it actually is the fact that this sometimes work and sometimes do not work if I set the route to /datasets. I'm suspecting it may be due to the way we do topological sort in our resolver

[ashwinb]

I don't think this is using cursor-based pagination because this is just an index we are returning? Maybe just avoid saying this is cursor-based?

[yanxi0830]

I'm basing the API off of this cursor based pagination spec.

[hardikjshah]

Where do we do validation ?

[yanxi0830]

We do validation at the eval API level:

• it will happen here:

  llama-stack/llama_stack/providers/inline/eval/meta_reference/eval.py

  Lines 87 to 88 in cc48d9e

  	# TODO (xiyan): validate dataset schema
  	# dataset_def = await self.datasets_api.get_dataset(dataset_id=dataset_id)

which will be added with the eval api implementation changes

[hardikjshah]

Should append_rows take source , similarly to dataset's registration ?
What is the use case where we need append ?

[yanxi0830]

I think @dineshyv added the append_rows API which is used internally for storing telemetry logs into a dataset.

But I don't feel like it is useful to be exposed to users, so we can remove it and make it an internal api in the future.

[hardikjshah]

Are we going to put all params in the uri ? or put them in uri_params ?
I know we went back and forth on this , but did we decide on the uri route ?

[yanxi0830]

We decided on the uri route based on @ashwinb's feedback on not having a separate HuggingfaceDataSource.