scripts/share_traces_to_hf_dataset.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import os
import subprocess
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

ROOT = Path(__file__).resolve().parents[1]
for candidate in (ROOT / 'src', ROOT):
    if candidate.is_dir() and str(candidate) not in sys.path:
        sys.path.insert(0, str(candidate))

from app_kit.tracing import canonicalize_trace_payload

SCHEMA_FIELDS = ['timestamp', 'inputs', 'parsed_outputs', 'model_name', 'model_id', 'adapter_name', 'generation_stats', 'checkpoint_path', 'checkpoint_source']


def _relative_or_absolute(path: Path, root: Path) -> str:
    try:
        return str(path.resolve().relative_to(root.resolve()))
    except Exception:
        return str(path.resolve())


def discover_trace_files(traces_root: str | Path) -> list[Path]:
    root = Path(traces_root).expanduser()
    if not root.exists():
        raise FileNotFoundError(f'Trace directory not found: {root}')
    if root.is_file():
        return [root]
    if root.name == 'traces':
        candidates = [path for path in sorted(root.glob('*.json')) if path.is_file()]
    else:
        candidates = [path for path in sorted(root.rglob('traces/*.json')) if path.is_file()]
    return candidates


def default_output_dir(repo_root: str | Path, repo_name: str | None = None, today: str | None = None) -> Path:
    root = Path(repo_root).expanduser()
    repo_name = repo_name or root.name
    today = today or datetime.now(timezone.utc).date().isoformat()
    return root / 'artifacts' / 'verification' / today / 'sharing_is_caring' / repo_name


def build_dataset_row(trace_file: Path, repo_root: Path, repo_name: str) -> dict[str, Any]:
    trace_payload = json.loads(trace_file.read_text(encoding='utf-8'))
    trace = canonicalize_trace_payload(trace_payload)
    row: dict[str, Any] = {
        'repo_name': repo_name,
        'source_trace_path': _relative_or_absolute(trace_file, repo_root),
        'timestamp': trace['timestamp'],
        'inputs': trace['inputs'],
        'parsed_outputs': trace['parsed_outputs'],
        'model_name': trace['model_name'],
        'kind': trace.get('kind', 'trace'),
    }
    for key in ('project', 'pack_id', 'pack_name', 'pack_path', 'model_id', 'adapter_name', 'checkpoint_path', 'checkpoint_source', 'generation_stats'):
        value = trace.get(key)
        if value not in (None, '', [], {}, ()):
            row[key] = value
    return row


def materialize_dataset(
    traces_dir: str | Path,
    output_dir: str | Path | None = None,
    repo_root: str | Path | None = None,
    repo_name: str | None = None,
    today: str | None = None,
) -> dict[str, Any]:
    traces_dir = Path(traces_dir).expanduser()
    repo_root = Path(repo_root).expanduser() if repo_root is not None else ROOT
    repo_name = repo_name or repo_root.name
    today = today or datetime.now(timezone.utc).date().isoformat()
    output_dir = Path(output_dir).expanduser() if output_dir is not None else default_output_dir(repo_root, repo_name, today)
    trace_files = discover_trace_files(traces_dir)
    if not trace_files:
        raise FileNotFoundError(f'No trace JSON files found under {traces_dir}')
    rows = [build_dataset_row(trace_file, repo_root=repo_root, repo_name=repo_name) for trace_file in trace_files]
    rows.sort(key=lambda row: row['source_trace_path'])
    output_dir.mkdir(parents=True, exist_ok=True)
    dataset_jsonl = output_dir / 'dataset.jsonl'
    metadata_json = output_dir / 'dataset-metadata.json'
    with dataset_jsonl.open('w', encoding='utf-8') as handle:
        for row in rows:
            handle.write(json.dumps(row, ensure_ascii=False, sort_keys=True) + '\n')
    metadata = {
        'repo_name': repo_name,
        'repo_root': str(repo_root),
        'source_traces_dir': str(traces_dir),
        'source_trace_count': len(trace_files),
        'records_written': len(trace_files),
        'source_trace_files': [row['source_trace_path'] for row in rows],
        'generated_date': today,
        'schema_fields': SCHEMA_FIELDS,
        'output_files': {
            'dataset_jsonl': dataset_jsonl.name,
            'metadata_json': metadata_json.name,
        },
    }
    metadata_json.write_text(json.dumps(metadata, indent=2, ensure_ascii=False, sort_keys=True), encoding='utf-8')
    return {
        'repo_name': repo_name,
        'repo_root': repo_root,
        'source_traces_dir': traces_dir,
        'trace_files': trace_files,
        'source_trace_count': len(trace_files),
        'records_written': len(trace_files),
        'output_dir': output_dir,
        'dataset_jsonl': dataset_jsonl,
        'metadata_json': metadata_json,
    }


def _run_git(args: list[str], cwd: Path) -> subprocess.CompletedProcess[str]:
    proc = subprocess.run(['git', *args], cwd=cwd, text=True, capture_output=True)
    if proc.returncode != 0:
        output = (proc.stdout + '\n' + proc.stderr).strip()
        raise RuntimeError(f"git {' '.join(args)} failed in {cwd}: {output}")
    return proc


def build_hf_git_url(repo_id: str, token: str) -> str:
    return f'https://__token__:{token}@huggingface.co/datasets/{repo_id}.git'


def git_push_dataset(output_dir: str | Path, remote_url: str, branch: str = 'main') -> dict[str, Any]:
    output_dir = Path(output_dir).expanduser()
    if not (output_dir / '.git').exists():
        _run_git(['init', '-b', branch], cwd=output_dir)
    _run_git(['config', 'user.name', 'Hermes Trace Sharer'], cwd=output_dir)
    _run_git(['config', 'user.email', 'trace-sharer@example.com'], cwd=output_dir)
    _run_git(['add', 'dataset.jsonl', 'dataset-metadata.json'], cwd=output_dir)
    commit = subprocess.run(
        ['git', 'commit', '-m', 'Share traces to Hugging Face Dataset'],
        cwd=output_dir,
        text=True,
        capture_output=True,
    )
    commit_output = (commit.stdout + '\n' + commit.stderr).strip().lower()
    if commit.returncode != 0 and 'nothing to commit' not in commit_output:
        raise RuntimeError(f'git commit failed in {output_dir}: {commit_output}')
    push = subprocess.run(
        ['git', 'push', remote_url, f'HEAD:{branch}'],
        cwd=output_dir,
        text=True,
        capture_output=True,
    )
    if push.returncode != 0:
        output = (push.stdout + '\n' + push.stderr).strip()
        raise RuntimeError(f'git push failed in {output_dir}: {output}')
    return {
        'pushed': True,
        'branch': branch,
    }


def push_to_hf_dataset(output_dir: str | Path, repo_id: str, token: str, branch: str = 'main') -> dict[str, Any]:
    return git_push_dataset(output_dir, build_hf_git_url(repo_id, token), branch=branch)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description='Materialize and optionally share traces as a Hugging Face Dataset bundle')
    parser.add_argument('traces_dir', nargs='?', default='.', help='Local trace directory or repo root to scan for trace JSON files')
    parser.add_argument('--output-dir', default=None, help='Optional output directory for the materialized dataset bundle')
    parser.add_argument('--repo-id', default=os.environ.get('HF_DATASET_REPO'), help='Hugging Face dataset repo id (namespace/name) for --push')
    parser.add_argument('--branch', default='main', help='Target branch when pushing')
    parser.add_argument('--push', action='store_true', help='Commit and push to Hugging Face when HF_TOKEN is present')
    parser.add_argument('--dry-run', action='store_true', help='Force local-only materialization even if --push is supplied')
    args = parser.parse_args(argv)

    materialized = materialize_dataset(args.traces_dir, output_dir=args.output_dir, repo_root=ROOT, repo_name=ROOT.name)
    push_status: dict[str, Any] = {
        'attempted': False,
        'pushed': False,
        'reason': 'local materialization only',
    }
    if args.push and not args.dry_run:
        repo_id = args.repo_id
        token = os.environ.get('HF_TOKEN') or os.environ.get('HUGGINGFACE_HUB_TOKEN')
        push_status['attempted'] = True
        if not repo_id:
            push_status['reason'] = 'HF dataset repo id missing; pass --repo-id or set HF_DATASET_REPO'
        elif not token:
            push_status['reason'] = 'HF_TOKEN missing; local materialization only'
        else:
            push_status = {
                'attempted': True,
                'pushed': True,
                'repo_id': repo_id,
                'branch': args.branch,
                **push_to_hf_dataset(materialized['output_dir'], repo_id, token, branch=args.branch),
            }
    elif args.push and args.dry_run:
        push_status = {
            'attempted': False,
            'pushed': False,
            'reason': 'dry-run requested',
        }

    summary = {
        'repo_name': materialized['repo_name'],
        'records_written': materialized['source_trace_count'],
        'output_dir': str(materialized['output_dir']),
        'dataset_jsonl': str(materialized['dataset_jsonl']),
        'metadata_json': str(materialized['metadata_json']),
        'trace_files': [str(path) for path in materialized['trace_files']],
        'push': push_status,
    }
    print(json.dumps(summary, indent=2, ensure_ascii=False, sort_keys=True))
    return 0


if __name__ == '__main__':
    raise SystemExit(main())