yourComplete/quickstart-trackio

1import os2from pathlib import Path3 4import pandas as pd5 6from trackio import deploy, utils7from trackio.sqlite_storage import SQLiteStorage8 9 10def import_csv(11    csv_path: str | Path,12    project: str,13    name: str | None = None,14    space_id: str | None = None,15    dataset_id: str | None = None,16    private: bool | None = None,17) -> None:18    """19    Imports a CSV file into a Trackio project. The CSV file must contain a `"step"`20    column, may optionally contain a `"timestamp"` column, and any other columns will be21    treated as metrics. It should also include a header row with the column names.22 23    TODO: call init() and return a Run object so that the user can continue to log metrics to it.24 25    Args:26        csv_path (`str` or `Path`):27            The str or Path to the CSV file to import.28        project (`str`):29            The name of the project to import the CSV file into. Must not be an existing30            project.31        name (`str`, *optional*):32            The name of the Run to import the CSV file into. If not provided, a default33            name will be generated.34        name (`str`, *optional*):35            The name of the run (if not provided, a default name will be generated).36        space_id (`str`, *optional*):37            If provided, the project will be logged to a Hugging Face Space instead of a38            local directory. Should be a complete Space name like `"username/reponame"`39            or `"orgname/reponame"`, or just `"reponame"` in which case the Space will40            be created in the currently-logged-in Hugging Face user's namespace. If the41            Space does not exist, it will be created. If the Space already exists, the42            project will be logged to it.43        dataset_id (`str`, *optional*):44            If provided, a persistent Hugging Face Dataset will be created and the45            metrics will be synced to it every 5 minutes. Should be a complete Dataset46            name like `"username/datasetname"` or `"orgname/datasetname"`, or just47            `"datasetname"` in which case the Dataset will be created in the48            currently-logged-in Hugging Face user's namespace. If the Dataset does not49            exist, it will be created. If the Dataset already exists, the project will50            be appended to it. If not provided, the metrics will be logged to a local51            SQLite database, unless a `space_id` is provided, in which case a Dataset52            will be automatically created with the same name as the Space but with the53            `"_dataset"` suffix.54        private (`bool`, *optional*):55            Whether to make the Space private. If None (default), the repo will be56            public unless the organization's default is private. This value is ignored57            if the repo already exists.58    """59    if SQLiteStorage.get_runs(project):60        raise ValueError(61            f"Project '{project}' already exists. Cannot import CSV into existing project."62        )63 64    csv_path = Path(csv_path)65    if not csv_path.exists():66        raise FileNotFoundError(f"CSV file not found: {csv_path}")67 68    df = pd.read_csv(csv_path)69    if df.empty:70        raise ValueError("CSV file is empty")71 72    column_mapping = utils.simplify_column_names(df.columns.tolist())73    df = df.rename(columns=column_mapping)74 75    step_column = None76    for col in df.columns:77        if col.lower() == "step":78            step_column = col79            break80 81    if step_column is None:82        raise ValueError("CSV file must contain a 'step' or 'Step' column")83 84    if name is None:85        name = csv_path.stem86 87    metrics_list = []88    steps = []89    timestamps = []90 91    numeric_columns = []92    for column in df.columns:93        if column == step_column:94            continue95        if column == "timestamp":96            continue97 98        try:99            pd.to_numeric(df[column], errors="raise")100            numeric_columns.append(column)101        except (ValueError, TypeError):102            continue103 104    for _, row in df.iterrows():105        metrics = {}106        for column in numeric_columns:107            value = row[column]108            if bool(pd.notna(value)):109                metrics[column] = float(value)110 111        if metrics:112            metrics_list.append(metrics)113            steps.append(int(row[step_column]))114 115            if "timestamp" in df.columns and bool(pd.notna(row["timestamp"])):116                timestamps.append(str(row["timestamp"]))117            else:118                timestamps.append("")119 120    if metrics_list:121        SQLiteStorage.bulk_log(122            project=project,123            run=name,124            metrics_list=metrics_list,125            steps=steps,126            timestamps=timestamps,127        )128 129    print(130        f"* Imported {len(metrics_list)} rows from {csv_path} into project '{project}' as run '{name}'"131    )132    print(f"* Metrics found: {', '.join(metrics_list[0].keys())}")133 134    space_id, dataset_id = utils.preprocess_space_and_dataset_ids(space_id, dataset_id)135    if dataset_id is not None:136        os.environ["TRACKIO_DATASET_ID"] = dataset_id137        print(f"* Trackio metrics will be synced to Hugging Face Dataset: {dataset_id}")138 139    if space_id is None:140        utils.print_dashboard_instructions(project)141    else:142        deploy.create_space_if_not_exists(143            space_id=space_id, dataset_id=dataset_id, private=private144        )145        deploy.wait_until_space_exists(space_id=space_id)146        deploy.upload_db_to_space(project=project, space_id=space_id)147        print(148            f"* View dashboard by going to: {deploy.SPACE_URL.format(space_id=space_id)}"149        )150 151 152def import_tf_events(153    log_dir: str | Path,154    project: str,155    name: str | None = None,156    space_id: str | None = None,157    dataset_id: str | None = None,158    private: bool | None = None,159) -> None:160    """161    Imports TensorFlow Events files from a directory into a Trackio project. Each162    subdirectory in the log directory will be imported as a separate run.163 164    Args:165        log_dir (`str` or `Path`):166            The str or Path to the directory containing TensorFlow Events files.167        project (`str`):168            The name of the project to import the TensorFlow Events files into. Must not169            be an existing project.170        name (`str`, *optional*):171            The name prefix for runs (if not provided, will use directory names). Each172            subdirectory will create a separate run.173        space_id (`str`, *optional*):174            If provided, the project will be logged to a Hugging Face Space instead of a175            local directory. Should be a complete Space name like `"username/reponame"`176            or `"orgname/reponame"`, or just `"reponame"` in which case the Space will177            be created in the currently-logged-in Hugging Face user's namespace. If the178            Space does not exist, it will be created. If the Space already exists, the179            project will be logged to it.180        dataset_id (`str`, *optional*):181            If provided, a persistent Hugging Face Dataset will be created and the182            metrics will be synced to it every 5 minutes. Should be a complete Dataset183            name like `"username/datasetname"` or `"orgname/datasetname"`, or just184            `"datasetname"` in which case the Dataset will be created in the185            currently-logged-in Hugging Face user's namespace. If the Dataset does not186            exist, it will be created. If the Dataset already exists, the project will187            be appended to it. If not provided, the metrics will be logged to a local188            SQLite database, unless a `space_id` is provided, in which case a Dataset189            will be automatically created with the same name as the Space but with the190            `"_dataset"` suffix.191        private (`bool`, *optional*):192            Whether to make the Space private. If None (default), the repo will be193            public unless the organization's default is private. This value is ignored194            if the repo already exists.195    """196    try:197        from tbparse import SummaryReader198    except ImportError:199        raise ImportError(200            "The `tbparse` package is not installed but is required for `import_tf_events`. Please install trackio with the `tensorboard` extra: `pip install trackio[tensorboard]`."201        )202 203    if SQLiteStorage.get_runs(project):204        raise ValueError(205            f"Project '{project}' already exists. Cannot import TF events into existing project."206        )207 208    path = Path(log_dir)209    if not path.exists():210        raise FileNotFoundError(f"TF events directory not found: {path}")211 212    # Use tbparse to read all tfevents files in the directory structure213    reader = SummaryReader(str(path), extra_columns={"dir_name"})214    df = reader.scalars215 216    if df.empty:217        raise ValueError(f"No TensorFlow events data found in {path}")218 219    total_imported = 0220    imported_runs = []221 222    # Group by dir_name to create separate runs223    for dir_name, group_df in df.groupby("dir_name"):224        try:225            # Determine run name based on directory name226            if dir_name == "":227                run_name = "main"  # For files in the root directory228            else:229                run_name = dir_name  # Use directory name230 231            if name:232                run_name = f"{name}_{run_name}"233 234            if group_df.empty:235                print(f"* Skipping directory {dir_name}: no scalar data found")236                continue237 238            metrics_list = []239            steps = []240            timestamps = []241 242            for _, row in group_df.iterrows():243                # Convert row values to appropriate types244                tag = str(row["tag"])245                value = float(row["value"])246                step = int(row["step"])247 248                metrics = {tag: value}249                metrics_list.append(metrics)250                steps.append(step)251 252                # Use wall_time if present, else fallback253                if "wall_time" in group_df.columns and not bool(254                    pd.isna(row["wall_time"])255                ):256                    timestamps.append(str(row["wall_time"]))257                else:258                    timestamps.append("")259 260            if metrics_list:261                SQLiteStorage.bulk_log(262                    project=project,263                    run=str(run_name),264                    metrics_list=metrics_list,265                    steps=steps,266                    timestamps=timestamps,267                )268 269                total_imported += len(metrics_list)270                imported_runs.append(run_name)271 272                print(273                    f"* Imported {len(metrics_list)} scalar events from directory '{dir_name}' as run '{run_name}'"274                )275                print(f"* Metrics in this run: {', '.join(set(group_df['tag']))}")276 277        except Exception as e:278            print(f"* Error processing directory {dir_name}: {e}")279            continue280 281    if not imported_runs:282        raise ValueError("No valid TensorFlow events data could be imported")283 284    print(f"* Total imported events: {total_imported}")285    print(f"* Created runs: {', '.join(imported_runs)}")286 287    space_id, dataset_id = utils.preprocess_space_and_dataset_ids(space_id, dataset_id)288    if dataset_id is not None:289        os.environ["TRACKIO_DATASET_ID"] = dataset_id290        print(f"* Trackio metrics will be synced to Hugging Face Dataset: {dataset_id}")291 292    if space_id is None:293        utils.print_dashboard_instructions(project)294    else:295        deploy.create_space_if_not_exists(296            space_id, dataset_id=dataset_id, private=private297        )298        deploy.wait_until_space_exists(space_id)299        deploy.upload_db_to_space(project, space_id)300        print(301            f"* View dashboard by going to: {deploy.SPACE_URL.format(space_id=space_id)}"302        )303