sync current ml-intern source for hf jobs

21b75e3 verified 5 months ago

12.7 kB

"""HF Hub artifact sync helpers for run directories and ETL caches."""

from __future__ import annotations

import argparse
from pathlib import Path

from huggingface_hub import HfApi, snapshot_download


# ---------------------------------------------------------------------------
# Model-repo helpers — run outputs (checkpoints, metrics, embeddings)
# ---------------------------------------------------------------------------


def push(repo: str, preset: str, source: Path, message: str = "snapshot") -> None:
    """Push one preset/run directory into the artifact repo namespace."""

    api = HfApi()
    api.create_repo(repo, repo_type="model", exist_ok=True, private=True)
    folder_path = source / preset
    if not folder_path.exists():
        raise FileNotFoundError(f"Missing preset artifact directory: {folder_path}")
    api.upload_folder(
        repo_id=repo,
        repo_type="model",
        folder_path=str(folder_path),
        path_in_repo=f"runs/{preset}",
        commit_message=message,
        ignore_patterns=["*.lmdb", "*.lmdb-lock"],
    )


def pull(repo: str, preset: str, dest: Path) -> None:
    """Pull one preset/run directory from the artifact repo."""

    dest.mkdir(parents=True, exist_ok=True)
    snapshot_download(
        repo_id=repo,
        repo_type="model",
        local_dir=str(dest),
        allow_patterns=[f"runs/{preset}/**"],
    )
    source = dest / "runs" / preset
    if not source.exists():
        return
    final_dir = dest / preset
    final_dir.mkdir(parents=True, exist_ok=True)
    for child in source.iterdir():
        target = final_dir / child.name
        if target.exists():
            if target.is_dir():
                continue
            target.unlink()
        child.rename(target)


# ---------------------------------------------------------------------------
# Dataset-repo helpers — ETL caches (LMDB, registry, classes) keyed by
# processed_signature.  Closes the gap documented in CLAUDE.md §6
# "Active gaps in HF coverage" by promoting the LMDB out of the model repo.
# ---------------------------------------------------------------------------


_ETL_ARTIFACTS = ("tensors.lmdb", "registry.csv", "classes.json", "etl_cache_manifest.json")
_QUANT_ARTIFACTS = (
    "transport_plan.pt",
    "alignment_offsets.pt",
    "tile_repo.pt",
    "read_data.pt",
    "counts.csv",
    "state.json",
    "summary.json",
)


def _etl_path_in_repo(processed_signature: str) -> str:
    return f"lmdb/{processed_signature}"


def sync_etl_to_dataset(
    repo: str,
    processed_signature: str,
    processed_root: Path,
    *,
    message: str = "etl-cache",
) -> dict[str, str]:
    """Upload an ETL output (`tensors.lmdb`, `registry.csv`, `classes.json`) to a dataset repo.
    `processed_signature` is the canonical cache key emitted by
    `deepgenopix matrix --json` (e.g. `px12_ps12_fl500-500_vf0p1_tf0p1_seed42_mf1`).
    The function uploads each present file individually so partial caches still
    push and tail files do not block earlier wins.
    """

    api = HfApi()
    api.create_repo(repo, repo_type="dataset", exist_ok=True, private=True)
    processed_root = Path(processed_root)
    uploaded: dict[str, str] = {}
    path_prefix = _etl_path_in_repo(processed_signature)
    for artifact in _ETL_ARTIFACTS:
        candidate = processed_root / artifact
        if not candidate.exists():
            continue
        api.upload_file(
            repo_id=repo,
            repo_type="dataset",
            path_or_fileobj=str(candidate),
            path_in_repo=f"{path_prefix}/{artifact}",
            commit_message=f"{message}: {artifact}",
        )
        uploaded[artifact] = f"{path_prefix}/{artifact}"
    return uploaded


def pull_etl_from_dataset(
    repo: str,
    processed_signature: str,
    processed_root: Path,
) -> dict[str, str]:
    """Pull the ETL cache for `processed_signature` from the dataset repo, if present.
    Returns a dict of `{artifact_name: local_path}` for each artifact that was
    materialized locally. An empty dict signals a cache miss — the caller is
    expected to run `deepgenopix etl` in that case.
    """

    processed_root = Path(processed_root)
    processed_root.mkdir(parents=True, exist_ok=True)
    path_prefix = _etl_path_in_repo(processed_signature)
    snapshot_dir = processed_root / "_hf_pull_cache"
    snapshot_dir.mkdir(parents=True, exist_ok=True)
    try:
        snapshot_download(
            repo_id=repo,
            repo_type="dataset",
            local_dir=str(snapshot_dir),
            allow_patterns=[f"{path_prefix}/**"],
        )
    except Exception:
        return {}

    pulled: dict[str, str] = {}
    source = snapshot_dir / "lmdb" / processed_signature
    if not source.exists():
        return pulled
    for artifact in _ETL_ARTIFACTS:
        candidate = source / artifact
        if not candidate.exists():
            continue
        destination = processed_root / artifact
        if destination.exists():
            destination.unlink()
        candidate.rename(destination)
        pulled[artifact] = str(destination)
    return pulled


# ---------------------------------------------------------------------------
# Dataset-repo helpers — quantification outputs consumed by variant jobs.
# These artifacts are large run inputs, not model checkpoints, so the canonical
# HF Jobs path keeps them in DATA_DATASET/quant_runs/<output_name>/.
# ---------------------------------------------------------------------------


def _quant_path_in_repo(output_name: str) -> str:
    return f"quant_runs/{output_name}"


def sync_quant_to_dataset(
    repo: str,
    output_name: str,
    quant_run_dir: Path,
    *,
    message: str = "quant-run",
) -> dict[str, str]:
    """Upload quantification artifacts into the dataset repo namespace."""

    api = HfApi()
    api.create_repo(repo, repo_type="dataset", exist_ok=True, private=True)
    quant_run_dir = Path(quant_run_dir)
    uploaded: dict[str, str] = {}
    path_prefix = _quant_path_in_repo(output_name)
    for artifact in _QUANT_ARTIFACTS:
        candidate = quant_run_dir / artifact
        if not candidate.exists():
            continue
        api.upload_file(
            repo_id=repo,
            repo_type="dataset",
            path_or_fileobj=str(candidate),
            path_in_repo=f"{path_prefix}/{artifact}",
            commit_message=f"{message}: {artifact}",
        )
        uploaded[artifact] = f"{path_prefix}/{artifact}"
    return uploaded


def pull_quant_from_dataset(
    repo: str,
    output_name: str,
    quant_run_dir: Path,
) -> dict[str, str]:
    """Pull quantification artifacts for `output_name` from the dataset repo."""

    quant_run_dir = Path(quant_run_dir)
    quant_run_dir.mkdir(parents=True, exist_ok=True)
    path_prefix = _quant_path_in_repo(output_name)
    snapshot_dir = quant_run_dir / "_hf_pull_cache"
    snapshot_dir.mkdir(parents=True, exist_ok=True)
    try:
        snapshot_download(
            repo_id=repo,
            repo_type="dataset",
            local_dir=str(snapshot_dir),
            allow_patterns=[f"{path_prefix}/**"],
        )
    except Exception:
        return {}

    pulled: dict[str, str] = {}
    source = snapshot_dir / "quant_runs" / output_name
    if not source.exists():
        return pulled
    for artifact in _QUANT_ARTIFACTS:
        candidate = source / artifact
        if not candidate.exists():
            continue
        destination = quant_run_dir / artifact
        if destination.exists():
            destination.unlink()
        candidate.rename(destination)
        pulled[artifact] = str(destination)
    return pulled


def push_dataset_folder(
    repo: str,
    source: Path,
    path_in_repo: str,
    *,
    message: str = "dataset-artifacts",
) -> None:
    """Upload a local folder into a dataset repo.
    This is used by Lane B jobs whose artifact shapes differ by corpus or sweep
    but all belong in dataset namespaces rather than the model artifact repo.
    """

    source = Path(source)
    if not source.exists():
        raise FileNotFoundError(f"Missing dataset artifact folder: {source}")
    api = HfApi()
    api.create_repo(repo, repo_type="dataset", exist_ok=True, private=True)
    api.upload_folder(
        repo_id=repo,
        repo_type="dataset",
        folder_path=str(source),
        path_in_repo=path_in_repo.strip("/"),
        commit_message=message,
    )


def main() -> int:
    parser = argparse.ArgumentParser(description="Push/pull DeepGenopix artifacts via the HF Hub")
    subparsers = parser.add_subparsers(dest="action", required=True)

    run_push = subparsers.add_parser("push", help="Push a run folder into the model repo")
    run_push.add_argument("--repo", required=True)
    run_push.add_argument("--preset", required=True)
    run_push.add_argument("--source", type=Path, default=Path("data/output"))
    run_push.add_argument("--message", default="snapshot")

    run_pull = subparsers.add_parser("pull", help="Pull a run folder from the model repo")
    run_pull.add_argument("--repo", required=True)
    run_pull.add_argument("--preset", required=True)
    run_pull.add_argument("--dest", type=Path, default=Path("data/output"))

    etl_push = subparsers.add_parser("push-etl", help="Push an ETL cache into a dataset repo")
    etl_push.add_argument("--repo", required=True)
    etl_push.add_argument("--signature", required=True)
    etl_push.add_argument("--processed-root", type=Path, required=True)
    etl_push.add_argument("--message", default="etl-cache")

    etl_pull = subparsers.add_parser("pull-etl", help="Pull an ETL cache from a dataset repo")
    etl_pull.add_argument("--repo", required=True)
    etl_pull.add_argument("--signature", required=True)
    etl_pull.add_argument("--processed-root", type=Path, required=True)

    quant_push = subparsers.add_parser("push-quant", help="Push quant artifacts into a dataset repo")
    quant_push.add_argument("--repo", required=True)
    quant_push.add_argument("--output-name", required=True)
    quant_push.add_argument("--quant-run-dir", type=Path, required=True)
    quant_push.add_argument("--message", default="quant-run")

    quant_pull = subparsers.add_parser("pull-quant", help="Pull quant artifacts from a dataset repo")
    quant_pull.add_argument("--repo", required=True)
    quant_pull.add_argument("--output-name", required=True)
    quant_pull.add_argument("--quant-run-dir", type=Path, required=True)

    dataset_folder_push = subparsers.add_parser(
        "push-dataset-folder", help="Push an arbitrary folder into a dataset repo"
    )
    dataset_folder_push.add_argument("--repo", required=True)
    dataset_folder_push.add_argument("--source", type=Path, required=True)
    dataset_folder_push.add_argument("--path-in-repo", required=True)
    dataset_folder_push.add_argument("--message", default="dataset-artifacts")

    args = parser.parse_args()
    if args.action == "push":
        push(args.repo, args.preset, args.source, args.message)
    elif args.action == "pull":
        pull(args.repo, args.preset, args.dest)
    elif args.action == "push-etl":
        result = sync_etl_to_dataset(
            args.repo, args.signature, args.processed_root, message=args.message
        )
        for artifact, remote_path in result.items():
            print(f"uploaded {artifact} → {remote_path}")
    elif args.action == "pull-etl":
        result = pull_etl_from_dataset(args.repo, args.signature, args.processed_root)
        if not result:
            print(f"cache miss for signature {args.signature}")
        else:
            for artifact, local_path in result.items():
                print(f"pulled {artifact} → {local_path}")
    elif args.action == "push-quant":
        result = sync_quant_to_dataset(
            args.repo, args.output_name, args.quant_run_dir, message=args.message
        )
        for artifact, remote_path in result.items():
            print(f"uploaded {artifact} → {remote_path}")
    elif args.action == "pull-quant":
        result = pull_quant_from_dataset(args.repo, args.output_name, args.quant_run_dir)
        if not result:
            print(f"cache miss for quant run {args.output_name}")
        else:
            for artifact, local_path in result.items():
                print(f"pulled {artifact} → {local_path}")
    elif args.action == "push-dataset-folder":
        push_dataset_folder(
            args.repo,
            args.source,
            args.path_in_repo,
            message=args.message,
        )
        print(f"uploaded {args.source} → {args.repo}/{args.path_in_repo.strip('/')}")
    return 0


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

	"""HF Hub artifact sync helpers for run directories and ETL caches."""

	from __future__ import annotations

	import argparse
	from pathlib import Path

	from huggingface_hub import HfApi, snapshot_download


	# ---------------------------------------------------------------------------
	# Model-repo helpers — run outputs (checkpoints, metrics, embeddings)
	# ---------------------------------------------------------------------------


	def push(repo: str, preset: str, source: Path, message: str = "snapshot") -> None:
	"""Push one preset/run directory into the artifact repo namespace."""

	api = HfApi()
	api.create_repo(repo, repo_type="model", exist_ok=True, private=True)
	folder_path = source / preset
	if not folder_path.exists():
	raise FileNotFoundError(f"Missing preset artifact directory: {folder_path}")
	api.upload_folder(
	repo_id=repo,
	repo_type="model",
	folder_path=str(folder_path),
	path_in_repo=f"runs/{preset}",
	commit_message=message,
	ignore_patterns=[".lmdb", ".lmdb-lock"],
	)


	def pull(repo: str, preset: str, dest: Path) -> None:
	"""Pull one preset/run directory from the artifact repo."""

	dest.mkdir(parents=True, exist_ok=True)
	snapshot_download(
	repo_id=repo,
	repo_type="model",
	local_dir=str(dest),
	allow_patterns=[f"runs/{preset}/**"],
	)
	source = dest / "runs" / preset
	if not source.exists():
	return
	final_dir = dest / preset
	final_dir.mkdir(parents=True, exist_ok=True)
	for child in source.iterdir():
	target = final_dir / child.name
	if target.exists():
	if target.is_dir():
	continue
	target.unlink()
	child.rename(target)


	# ---------------------------------------------------------------------------
	# Dataset-repo helpers — ETL caches (LMDB, registry, classes) keyed by
	# processed_signature. Closes the gap documented in CLAUDE.md §6
	# "Active gaps in HF coverage" by promoting the LMDB out of the model repo.
	# ---------------------------------------------------------------------------


	_ETL_ARTIFACTS = ("tensors.lmdb", "registry.csv", "classes.json", "etl_cache_manifest.json")
	_QUANT_ARTIFACTS = (
	"transport_plan.pt",
	"alignment_offsets.pt",
	"tile_repo.pt",
	"read_data.pt",
	"counts.csv",
	"state.json",
	"summary.json",
	)


	def _etl_path_in_repo(processed_signature: str) -> str:
	return f"lmdb/{processed_signature}"


	def sync_etl_to_dataset(
	repo: str,
	processed_signature: str,
	processed_root: Path,
	*,
	message: str = "etl-cache",
	) -> dict[str, str]:
	"""Upload an ETL output (`tensors.lmdb`, `registry.csv`, `classes.json`) to a dataset repo.

	`processed_signature` is the canonical cache key emitted by
	`deepgenopix matrix --json` (e.g. `px12_ps12_fl500-500_vf0p1_tf0p1_seed42_mf1`).
	The function uploads each present file individually so partial caches still
	push and tail files do not block earlier wins.
	"""

	api = HfApi()
	api.create_repo(repo, repo_type="dataset", exist_ok=True, private=True)
	processed_root = Path(processed_root)
	uploaded: dict[str, str] = {}
	path_prefix = _etl_path_in_repo(processed_signature)
	for artifact in _ETL_ARTIFACTS:
	candidate = processed_root / artifact
	if not candidate.exists():
	continue
	api.upload_file(
	repo_id=repo,
	repo_type="dataset",
	path_or_fileobj=str(candidate),
	path_in_repo=f"{path_prefix}/{artifact}",
	commit_message=f"{message}: {artifact}",
	)
	uploaded[artifact] = f"{path_prefix}/{artifact}"
	return uploaded


	def pull_etl_from_dataset(
	repo: str,
	processed_signature: str,
	processed_root: Path,
	) -> dict[str, str]:
	"""Pull the ETL cache for `processed_signature` from the dataset repo, if present.

	Returns a dict of `{artifact_name: local_path}` for each artifact that was
	materialized locally. An empty dict signals a cache miss — the caller is
	expected to run `deepgenopix etl` in that case.
	"""

	processed_root = Path(processed_root)
	processed_root.mkdir(parents=True, exist_ok=True)
	path_prefix = _etl_path_in_repo(processed_signature)
	snapshot_dir = processed_root / "_hf_pull_cache"
	snapshot_dir.mkdir(parents=True, exist_ok=True)
	try:
	snapshot_download(
	repo_id=repo,
	repo_type="dataset",
	local_dir=str(snapshot_dir),
	allow_patterns=[f"{path_prefix}/**"],
	)
	except Exception:
	return {}

	pulled: dict[str, str] = {}
	source = snapshot_dir / "lmdb" / processed_signature
	if not source.exists():
	return pulled
	for artifact in _ETL_ARTIFACTS:
	candidate = source / artifact
	if not candidate.exists():
	continue
	destination = processed_root / artifact
	if destination.exists():
	destination.unlink()
	candidate.rename(destination)
	pulled[artifact] = str(destination)
	return pulled


	# ---------------------------------------------------------------------------
	# Dataset-repo helpers — quantification outputs consumed by variant jobs.
	# These artifacts are large run inputs, not model checkpoints, so the canonical
	# HF Jobs path keeps them in DATA_DATASET/quant_runs/<output_name>/.
	# ---------------------------------------------------------------------------


	def _quant_path_in_repo(output_name: str) -> str:
	return f"quant_runs/{output_name}"


	def sync_quant_to_dataset(
	repo: str,
	output_name: str,
	quant_run_dir: Path,
	*,
	message: str = "quant-run",
	) -> dict[str, str]:
	"""Upload quantification artifacts into the dataset repo namespace."""

	api = HfApi()
	api.create_repo(repo, repo_type="dataset", exist_ok=True, private=True)
	quant_run_dir = Path(quant_run_dir)
	uploaded: dict[str, str] = {}
	path_prefix = _quant_path_in_repo(output_name)
	for artifact in _QUANT_ARTIFACTS:
	candidate = quant_run_dir / artifact
	if not candidate.exists():
	continue
	api.upload_file(
	repo_id=repo,
	repo_type="dataset",
	path_or_fileobj=str(candidate),
	path_in_repo=f"{path_prefix}/{artifact}",
	commit_message=f"{message}: {artifact}",
	)
	uploaded[artifact] = f"{path_prefix}/{artifact}"
	return uploaded


	def pull_quant_from_dataset(
	repo: str,
	output_name: str,
	quant_run_dir: Path,
	) -> dict[str, str]:
	"""Pull quantification artifacts for `output_name` from the dataset repo."""

	quant_run_dir = Path(quant_run_dir)
	quant_run_dir.mkdir(parents=True, exist_ok=True)
	path_prefix = _quant_path_in_repo(output_name)
	snapshot_dir = quant_run_dir / "_hf_pull_cache"
	snapshot_dir.mkdir(parents=True, exist_ok=True)
	try:
	snapshot_download(
	repo_id=repo,
	repo_type="dataset",
	local_dir=str(snapshot_dir),
	allow_patterns=[f"{path_prefix}/**"],
	)
	except Exception:
	return {}

	pulled: dict[str, str] = {}
	source = snapshot_dir / "quant_runs" / output_name
	if not source.exists():
	return pulled
	for artifact in _QUANT_ARTIFACTS:
	candidate = source / artifact
	if not candidate.exists():
	continue
	destination = quant_run_dir / artifact
	if destination.exists():
	destination.unlink()
	candidate.rename(destination)
	pulled[artifact] = str(destination)
	return pulled


	def push_dataset_folder(
	repo: str,
	source: Path,
	path_in_repo: str,
	*,
	message: str = "dataset-artifacts",
	) -> None:
	"""Upload a local folder into a dataset repo.

	This is used by Lane B jobs whose artifact shapes differ by corpus or sweep
	but all belong in dataset namespaces rather than the model artifact repo.
	"""

	source = Path(source)
	if not source.exists():
	raise FileNotFoundError(f"Missing dataset artifact folder: {source}")
	api = HfApi()
	api.create_repo(repo, repo_type="dataset", exist_ok=True, private=True)
	api.upload_folder(
	repo_id=repo,
	repo_type="dataset",
	folder_path=str(source),
	path_in_repo=path_in_repo.strip("/"),
	commit_message=message,
	)


	def main() -> int:
	parser = argparse.ArgumentParser(description="Push/pull DeepGenopix artifacts via the HF Hub")
	subparsers = parser.add_subparsers(dest="action", required=True)

	run_push = subparsers.add_parser("push", help="Push a run folder into the model repo")
	run_push.add_argument("--repo", required=True)
	run_push.add_argument("--preset", required=True)
	run_push.add_argument("--source", type=Path, default=Path("data/output"))
	run_push.add_argument("--message", default="snapshot")

	run_pull = subparsers.add_parser("pull", help="Pull a run folder from the model repo")
	run_pull.add_argument("--repo", required=True)
	run_pull.add_argument("--preset", required=True)
	run_pull.add_argument("--dest", type=Path, default=Path("data/output"))

	etl_push = subparsers.add_parser("push-etl", help="Push an ETL cache into a dataset repo")
	etl_push.add_argument("--repo", required=True)
	etl_push.add_argument("--signature", required=True)
	etl_push.add_argument("--processed-root", type=Path, required=True)
	etl_push.add_argument("--message", default="etl-cache")

	etl_pull = subparsers.add_parser("pull-etl", help="Pull an ETL cache from a dataset repo")
	etl_pull.add_argument("--repo", required=True)
	etl_pull.add_argument("--signature", required=True)
	etl_pull.add_argument("--processed-root", type=Path, required=True)

	quant_push = subparsers.add_parser("push-quant", help="Push quant artifacts into a dataset repo")
	quant_push.add_argument("--repo", required=True)
	quant_push.add_argument("--output-name", required=True)
	quant_push.add_argument("--quant-run-dir", type=Path, required=True)
	quant_push.add_argument("--message", default="quant-run")

	quant_pull = subparsers.add_parser("pull-quant", help="Pull quant artifacts from a dataset repo")
	quant_pull.add_argument("--repo", required=True)
	quant_pull.add_argument("--output-name", required=True)
	quant_pull.add_argument("--quant-run-dir", type=Path, required=True)

	dataset_folder_push = subparsers.add_parser(
	"push-dataset-folder", help="Push an arbitrary folder into a dataset repo"
	)
	dataset_folder_push.add_argument("--repo", required=True)
	dataset_folder_push.add_argument("--source", type=Path, required=True)
	dataset_folder_push.add_argument("--path-in-repo", required=True)
	dataset_folder_push.add_argument("--message", default="dataset-artifacts")

	args = parser.parse_args()
	if args.action == "push":
	push(args.repo, args.preset, args.source, args.message)
	elif args.action == "pull":
	pull(args.repo, args.preset, args.dest)
	elif args.action == "push-etl":
	result = sync_etl_to_dataset(
	args.repo, args.signature, args.processed_root, message=args.message
	)
	for artifact, remote_path in result.items():
	print(f"uploaded {artifact} → {remote_path}")
	elif args.action == "pull-etl":
	result = pull_etl_from_dataset(args.repo, args.signature, args.processed_root)
	if not result:
	print(f"cache miss for signature {args.signature}")
	else:
	for artifact, local_path in result.items():
	print(f"pulled {artifact} → {local_path}")
	elif args.action == "push-quant":
	result = sync_quant_to_dataset(
	args.repo, args.output_name, args.quant_run_dir, message=args.message
	)
	for artifact, remote_path in result.items():
	print(f"uploaded {artifact} → {remote_path}")
	elif args.action == "pull-quant":
	result = pull_quant_from_dataset(args.repo, args.output_name, args.quant_run_dir)
	if not result:
	print(f"cache miss for quant run {args.output_name}")
	else:
	for artifact, local_path in result.items():
	print(f"pulled {artifact} → {local_path}")
	elif args.action == "push-dataset-folder":
	push_dataset_folder(
	args.repo,
	args.source,
	args.path_in_repo,
	message=args.message,
	)
	print(f"uploaded {args.source} → {args.repo}/{args.path_in_repo.strip('/')}")
	return 0


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