wehub-resource-sync
139 lines
3.8 KiB
Python
139 lines
3.8 KiB
Python
import json
|
|
from typing import Any, Literal
|
|
|
|
import click
|
|
|
|
from mlflow import MlflowClient
|
|
from mlflow.environment_variables import MLFLOW_EXPERIMENT_ID
|
|
from mlflow.utils.string_utils import _create_table
|
|
from mlflow.utils.time import conv_longdate_to_str
|
|
|
|
EXPERIMENT_ID = click.option(
|
|
"--experiment-id",
|
|
"-x",
|
|
envvar=MLFLOW_EXPERIMENT_ID.name,
|
|
type=click.STRING,
|
|
required=True,
|
|
help="Experiment ID to list datasets for. Can be set via MLFLOW_EXPERIMENT_ID env var.",
|
|
)
|
|
|
|
|
|
def _format_datasets_as_json(datasets) -> dict[str, Any]:
|
|
"""Format datasets as a JSON-serializable dictionary."""
|
|
return {
|
|
"datasets": [
|
|
{
|
|
"dataset_id": ds.dataset_id,
|
|
"name": ds.name,
|
|
"digest": ds.digest,
|
|
"created_time": ds.created_time,
|
|
"last_update_time": ds.last_update_time,
|
|
"created_by": ds.created_by,
|
|
"last_updated_by": ds.last_updated_by,
|
|
"tags": ds.tags,
|
|
}
|
|
for ds in datasets
|
|
],
|
|
"next_page_token": datasets.token,
|
|
}
|
|
|
|
|
|
def _format_datasets_as_table(datasets) -> tuple[list[list[str]], list[str]]:
|
|
"""Format datasets as table rows with headers."""
|
|
headers = ["Dataset ID", "Name", "Created", "Last Updated", "Created By"]
|
|
rows = []
|
|
for ds in datasets:
|
|
created = conv_longdate_to_str(ds.created_time) if ds.created_time else ""
|
|
updated = conv_longdate_to_str(ds.last_update_time) if ds.last_update_time else ""
|
|
rows.append([ds.dataset_id, ds.name, created, updated, ds.created_by or ""])
|
|
return rows, headers
|
|
|
|
|
|
@click.group("datasets")
|
|
def commands():
|
|
"""Manage GenAI evaluation datasets."""
|
|
|
|
|
|
@commands.command("list")
|
|
@EXPERIMENT_ID
|
|
@click.option(
|
|
"--filter-string",
|
|
type=click.STRING,
|
|
help="Filter string (e.g., \"name LIKE 'qa_%'\").",
|
|
)
|
|
@click.option(
|
|
"--max-results",
|
|
type=click.INT,
|
|
default=50,
|
|
help="Maximum results (default: 50).",
|
|
)
|
|
@click.option(
|
|
"--order-by",
|
|
type=click.STRING,
|
|
help="Columns to order by (e.g., 'last_update_time DESC').",
|
|
)
|
|
@click.option(
|
|
"--page-token",
|
|
type=click.STRING,
|
|
help="Pagination token.",
|
|
)
|
|
@click.option(
|
|
"--output",
|
|
type=click.Choice(["table", "json"]),
|
|
default="table",
|
|
help="Output format.",
|
|
)
|
|
def list_datasets(
|
|
experiment_id: str,
|
|
filter_string: str | None = None,
|
|
max_results: int = 50,
|
|
order_by: str | None = None,
|
|
page_token: str | None = None,
|
|
output: Literal["table", "json"] = "table",
|
|
) -> None:
|
|
"""
|
|
List GenAI evaluation datasets associated with an experiment.
|
|
|
|
\b
|
|
Examples:
|
|
# List datasets in experiment 1
|
|
mlflow datasets list --experiment-id 1
|
|
|
|
\b
|
|
# Using environment variable
|
|
export MLFLOW_EXPERIMENT_ID=1
|
|
mlflow datasets list --max-results 10
|
|
|
|
\b
|
|
# Filter datasets by name pattern
|
|
mlflow datasets list --experiment-id 1 --filter-string "name LIKE 'qa_%'"
|
|
|
|
\b
|
|
# Order results by last update time
|
|
mlflow datasets list --experiment-id 1 --order-by "last_update_time DESC"
|
|
|
|
\b
|
|
# Output as JSON
|
|
mlflow datasets list --experiment-id 1 --output json
|
|
"""
|
|
client = MlflowClient()
|
|
order_by_list = [o.strip() for o in order_by.split(",")] if order_by else None
|
|
|
|
datasets = client.search_datasets(
|
|
experiment_ids=[experiment_id],
|
|
filter_string=filter_string,
|
|
max_results=max_results,
|
|
order_by=order_by_list,
|
|
page_token=page_token,
|
|
)
|
|
|
|
if output == "json":
|
|
result = _format_datasets_as_json(datasets)
|
|
click.echo(json.dumps(result, indent=2))
|
|
else:
|
|
rows, headers = _format_datasets_as_table(datasets)
|
|
click.echo(_create_table(rows, headers=headers))
|
|
|
|
if datasets.token:
|
|
click.echo(f"\nNext page token: {datasets.token}")
|