| import os |
| import json |
| import time |
| import logging |
| from pathlib import Path |
| from typing import Any, Dict, List, Optional, Tuple |
|
|
| import torch |
| from transformers import Trainer |
|
|
| from .manifest import build_manifest_from_dir |
| from .vocab import build_gene_vocab_from_manifest, build_meta_vocab_from_manifest |
|
|
|
|
| logger = logging.getLogger("scRep_train") |
|
|
|
|
| |
| |
| |
| def dist_is_initialized() -> bool: |
| return torch.distributed.is_available() and torch.distributed.is_initialized() |
|
|
|
|
| def dist_rank_world() -> Tuple[int, int]: |
| if dist_is_initialized(): |
| return torch.distributed.get_rank(), torch.distributed.get_world_size() |
| rank = int(os.environ.get("RANK", "0")) |
| world = int(os.environ.get("WORLD_SIZE", "1")) |
| return rank, world |
|
|
|
|
| def dist_barrier() -> None: |
| if dist_is_initialized(): |
| torch.distributed.barrier() |
|
|
|
|
| def is_world_process_zero() -> bool: |
| rank, _ = dist_rank_world() |
| return rank == 0 |
|
|
|
|
| def wait_for_file(path: str, timeout: float = 300.0, poll: float = 0.5) -> None: |
| """Wait until a file exists and is readable.""" |
| p = Path(path) |
| start = time.time() |
| last_err = None |
| while time.time() - start < timeout: |
| if p.exists(): |
| try: |
| with open(p, "r") as f: |
| _ = f.read(1) |
| return |
| except Exception as e: |
| last_err = e |
| time.sleep(poll) |
| raise TimeoutError(f"Timed out waiting for file: {path}. Last error: {last_err}") |
|
|
|
|
| |
| |
| |
| def load_json(path: str) -> Any: |
| with open(path, "r") as f: |
| return json.load(f) |
|
|
|
|
| def save_json(obj: Any, path: str) -> None: |
| os.makedirs(os.path.dirname(path), exist_ok=True) |
| tmp_path = path + ".tmp" |
| with open(tmp_path, "w") as f: |
| json.dump(obj, f, ensure_ascii=False, indent=2) |
| f.flush() |
| os.fsync(f.fileno()) |
| os.replace(tmp_path, path) |
|
|
|
|
| |
| |
| |
| def parse_fsdp_list(fsdp_str: Optional[str]) -> Optional[List[str]]: |
| s = (fsdp_str or "").strip() |
| if not s: |
| return None |
| s = s.replace(",", " ") |
| parts = [p.strip() for p in s.split() if p.strip()] |
| return parts if parts else None |
|
|
|
|
| def load_fsdp_config(path: Optional[str]) -> Optional[Dict[str, Any]]: |
| if not path: |
| return None |
| if not os.path.exists(path): |
| raise FileNotFoundError(f"fsdp_config not found: {path}") |
| return load_json(path) |
|
|
|
|
| def normalize_report_to(report_to: str): |
| return [] if report_to == "none" else report_to |
|
|
|
|
| class TrainerBatchWrapperCollator: |
| """Wrap the collated batch as {"batch": ...} for Trainer.""" |
| def __init__(self, base_collator): |
| self.base = base_collator |
|
|
| def __call__(self, examples): |
| return {"batch": self.base(examples)} |
|
|
|
|
| |
| |
| |
| def gene_vocab_to_map(gene_vocab: Dict[str, Any]) -> Dict[str, int]: |
| gene_offset = int(gene_vocab.get("gene_offset", 2)) |
| gene_names = list(gene_vocab["gene_names"]) |
| return {str(name): i + gene_offset for i, name in enumerate(gene_names)} |
|
|
|
|
| def _normalize_meta_keys(meta_keys: str) -> List[str]: |
| return [k.strip() for k in meta_keys.split(",") if k.strip()] |
|
|
|
|
| def _asset_metadata(args) -> Dict[str, Any]: |
| return { |
| "data_dir": os.path.abspath(args.data_dir), |
| "manifest_json": os.path.abspath(args.manifest_json) if args.manifest_json else "", |
| "gene_vocab_json": os.path.abspath(args.gene_vocab_json) if getattr(args, "gene_vocab_json", "") else "", |
| "use_raw": bool(args.use_raw), |
| "backed": args.backed, |
| "meta_keys": _normalize_meta_keys(args.meta_keys), |
| } |
|
|
|
|
| def prepare_scratch_assets(args, output_dir: str) -> Tuple[Any, Dict[str, Any], Dict[str, Any], Dict[str, int], Dict[str, Dict[str, int]], Dict[str, int]]: |
| is_main = is_world_process_zero() |
|
|
| manifest_path = os.path.join(output_dir, "manifest.json") |
| gene_vocab_path = os.path.join(output_dir, "gene_vocab.json") |
| meta_vocab_path = os.path.join(output_dir, "meta_vocab.json") |
| metadata_path = os.path.join(output_dir, "asset_metadata.json") |
|
|
| expected_metadata = _asset_metadata(args) |
| cache_complete = all( |
| os.path.exists(path) |
| for path in (manifest_path, gene_vocab_path, meta_vocab_path, metadata_path) |
| ) |
| metadata_matches = False |
| cache_reason = "missing cached asset files" |
| if cache_complete: |
| cached_metadata = load_json(metadata_path) |
| metadata_matches = cached_metadata == expected_metadata |
| if not metadata_matches: |
| cache_reason = "cached asset metadata does not match current dataset/config" |
|
|
| if cache_complete and metadata_matches: |
| logger.info("loading cached scratch assets from %s", output_dir) |
| manifest = load_json(manifest_path) |
| gene_vocab = load_json(gene_vocab_path) |
| meta_vocab = load_json(meta_vocab_path) |
| gene_name_to_id = gene_vocab_to_map(gene_vocab) |
| meta_maps = meta_vocab["meta_maps"] |
| meta_vocab_sizes = meta_vocab["meta_vocab_sizes"] |
| logger.info( |
| "loaded cached scratch assets: manifest=%s datasets, genes=%s, meta_keys=%s", |
| len(manifest), |
| gene_vocab.get("gene_vocab_size", 0), |
| list(meta_vocab_sizes.keys()), |
| ) |
| return manifest, gene_vocab, meta_vocab, gene_name_to_id, meta_maps, meta_vocab_sizes |
|
|
| if is_main: |
| logger.info("scratch assets cache miss in %s: %s", output_dir, cache_reason) |
| os.makedirs(output_dir, exist_ok=True) |
|
|
| if args.manifest_json and os.path.exists(args.manifest_json): |
| logger.info("building manifest from explicit manifest_json=%s", args.manifest_json) |
| manifest = load_json(args.manifest_json) |
| save_json(manifest, manifest_path) |
| else: |
| logger.info("building manifest from data_dir=%s", args.data_dir) |
| manifest_start = time.time() |
| manifest = build_manifest_from_dir( |
| args.data_dir, |
| backed=args.backed, |
| use_raw=bool(args.use_raw), |
| ) |
| if len(manifest) == 0: |
| raise RuntimeError( |
| "manifest is empty when building from `data_dir`. " |
| "This usually means your h5ad files were filtered out.\n\n" |
| "Common causes:\n" |
| "1) `--use_raw` is enabled but your .h5ad has no `adata.raw`.\n" |
| "2) Wrong `data_dir` / file pattern (no .h5ad found).\n" |
| "3) Required obs/meta columns mismatch (meta_keys).\n\n" |
| f"args.use_raw={args.use_raw}, data_dir={args.data_dir}" |
| ) |
| save_json(manifest, manifest_path) |
| logger.info("manifest saved to %s in %.1fs", manifest_path, time.time() - manifest_start) |
|
|
| meta_keys = _normalize_meta_keys(args.meta_keys) |
| if getattr(args, "gene_vocab_json", ""): |
| logger.info("loading fixed gene vocab from %s", args.gene_vocab_json) |
| gene_vocab = load_json(args.gene_vocab_json) |
| gene_names = list(gene_vocab["gene_names"]) |
| gene_name_to_id = gene_vocab_to_map(gene_vocab) |
| logger.info("fixed gene vocab loaded with %s genes", len(gene_names)) |
| else: |
| logger.info("building gene vocab from %s manifest entries", len(manifest)) |
| gene_start = time.time() |
| gene_name_to_id, gene_names = build_gene_vocab_from_manifest( |
| manifest, |
| use_raw=args.use_raw, |
| backed=args.backed, |
| gene_offset=2, |
| sort_paths=True, |
| ) |
| logger.info("gene vocab built with %s genes in %.1fs", len(gene_names), time.time() - gene_start) |
|
|
| logger.info("building meta vocab for keys=%s", meta_keys) |
| meta_start = time.time() |
| meta_maps, meta_vocab_sizes = build_meta_vocab_from_manifest( |
| manifest, |
| keys=meta_keys, |
| backed=args.backed, |
| sort_paths=True, |
| ) |
| logger.info("meta vocab built in %.1fs", time.time() - meta_start) |
|
|
| if not getattr(args, "gene_vocab_json", ""): |
| gene_vocab = { |
| "gene_offset": 2, |
| "gene_names": gene_names, |
| "gene_vocab_size": max(gene_name_to_id.values()) + 1, |
| } |
| meta_vocab = { |
| "meta_vocab_sizes": meta_vocab_sizes, |
| "meta_maps": meta_maps, |
| } |
| save_json(gene_vocab, gene_vocab_path) |
| save_json(meta_vocab, meta_vocab_path) |
| save_json(expected_metadata, metadata_path) |
| logger.info("scratch assets saved to %s", output_dir) |
| else: |
| logger.info("waiting for scratch assets to be prepared in %s", output_dir) |
| wait_for_file(metadata_path) |
| manifest = load_json(manifest_path) |
| gene_vocab = load_json(gene_vocab_path) |
| meta_vocab = load_json(meta_vocab_path) |
| gene_name_to_id = gene_vocab_to_map(gene_vocab) |
| meta_maps = meta_vocab["meta_maps"] |
| meta_vocab_sizes = meta_vocab["meta_vocab_sizes"] |
| logger.info("loaded scratch assets prepared by rank0 from %s", output_dir) |
|
|
| return manifest, gene_vocab, meta_vocab, gene_name_to_id, meta_maps, meta_vocab_sizes |
|
|
|
|
| def save_bundle( |
| trainer: Trainer, |
| model, |
| save_dir: str, |
| *, |
| gene_vocab: Dict[str, Any], |
| meta_vocab: Dict[str, Any], |
| manifest: Any, |
| ) -> None: |
| """ |
| Save a final bundle that is directly loadable by ``CellContrastivePredictModel.from_pretrained``. |
| |
| Under FSDP we explicitly ask Accelerate for a gathered state dict, then save it |
| through the unwrapped model. All ranks participate in ``get_state_dict``; only |
| world-process-zero writes files. |
| """ |
| os.makedirs(save_dir, exist_ok=True) |
|
|
| if hasattr(trainer, "accelerator") and trainer.accelerator is not None: |
| unwrapped_model = trainer.accelerator.unwrap_model(model) |
| state_dict = trainer.accelerator.get_state_dict(trainer.model) |
| if trainer.is_world_process_zero(): |
| unwrapped_model.save_pretrained( |
| save_dir, |
| save_safetensors=bool(trainer.args.save_safetensors), |
| state_dict=state_dict, |
| gene_vocab=gene_vocab, |
| meta_vocab=meta_vocab, |
| manifest=manifest, |
| extra_files={"training_args.json": trainer.args.to_dict()}, |
| ) |
| else: |
| if not trainer.is_world_process_zero(): |
| return |
| model.save_pretrained( |
| save_dir, |
| save_safetensors=bool(trainer.args.save_safetensors), |
| gene_vocab=gene_vocab, |
| meta_vocab=meta_vocab, |
| manifest=manifest, |
| extra_files={"training_args.json": trainer.args.to_dict()}, |
| ) |
|
|
