training/check_env.py

#!/usr/bin/env python
"""
Environment check for cpt_gemma4_sinhala.py (Unsloth + Gemma 4 E4B on an A40).

    python check_env.py            # quick checks (no model download)
    python check_env.py --model    # also loads Gemma 4 E4B, adds LoRA, runs 1 train step (~16 GB download)

Exit code 0 = everything passed.
"""
import argparse
import importlib
import importlib.metadata as md
import subprocess
import sys

OK, BAD, WARN = "[ OK ]", "[FAIL]", "[WARN]"
failures = []


def report(tag, msg):
    print(f"{tag} {msg}")
    if tag == BAD:
        failures.append(msg)


def ver(pkg):
    try:
        return md.version(pkg)
    except md.PackageNotFoundError:
        return None


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", action="store_true", help="load model + 1 training step")
    ap.add_argument("--model_name", default="unsloth/gemma-4-E4B")
    args = ap.parse_args()

    print(f"Python {sys.version.split()[0]}")

    # 1. GPU / torch  (import unsloth FIRST, as in the training script)
    try:
        import unsloth  # noqa: F401
        report(OK, f"unsloth imported (v{ver('unsloth')})")
    except Exception as e:  # noqa: BLE001
        report(BAD, f"unsloth import failed: {type(e).__name__}: {e}")

    try:
        import torch
        report(OK, f"torch {torch.__version__}, CUDA build {torch.version.cuda}")
        if torch.cuda.is_available():
            p = torch.cuda.get_device_properties(0)
            report(OK, f"GPU: {p.name}, {p.total_memory / 1024**3:.1f} GB, bf16={torch.cuda.is_bf16_supported()}")
        else:
            report(BAD, "torch.cuda.is_available() is False (Unsloth needs a GPU)")
    except Exception as e:  # noqa: BLE001
        report(BAD, f"torch import failed: {e}")

    # 2. Pinned versions
    transformers_v = ver("transformers")
    if transformers_v == "5.5.0":
        report(OK, "transformers 5.5.0")
    else:
        report(BAD, f"transformers is {transformers_v}, expected 5.5.0 "
                    "(re-run: pip install --no-deps transformers==5.5.0)")

    tok_v = ver("tokenizers")
    try:
        from packaging.version import Version
        good = tok_v and Version("0.22.0") <= Version(tok_v) <= Version("0.23.0")
    except Exception:  # noqa: BLE001
        good = False
    report(OK if good else BAD, f"tokenizers {tok_v} (need >=0.22.0,<=0.23.0)")

    hub_v = ver("huggingface_hub")
    try:
        from packaging.version import Version
        good = hub_v and Version("1.5.0") <= Version(hub_v) < Version("2.0")
    except Exception:  # noqa: BLE001
        good = False
    report(OK if good else BAD, f"huggingface_hub {hub_v} (need >=1.5.0,<2.0)")

    # 3. Other packages the script imports
    for pkg in ["trl", "peft", "datasets", "pandas", "accelerate", "bitsandbytes", "timm", "hf_transfer"]:
        v = ver(pkg)
        report(OK if v else BAD, f"{pkg} {v or 'NOT INSTALLED'}")

    # 4. Imports the training script actually does
    for stmt in ["from unsloth import FastModel",
                 "from trl import SFTConfig, SFTTrainer",
                 "from datasets import Dataset",
                 "from huggingface_hub import hf_hub_download"]:
        try:
            exec(stmt, {})  # noqa: S102
            report(OK, stmt)
        except Exception as e:  # noqa: BLE001
            report(BAD, f"{stmt} -> {type(e).__name__}: {e}")

    # 5. Do the SFTConfig args used by the script exist in this TRL version?
    try:
        from trl import SFTConfig
        import inspect
        fields = set(inspect.signature(SFTConfig.__init__).parameters)
        for a in ["max_length", "packing", "dataset_text_field", "bf16", "optim",
                  "eval_strategy", "warmup_ratio"]:
            report(OK if a in fields else BAD, f"SFTConfig accepts '{a}'")
    except Exception as e:  # noqa: BLE001
        report(WARN, f"could not inspect SFTConfig: {e}")

    # 6. pip's own dependency check (informational: --no-deps installs can show conflicts)
    out = subprocess.run([sys.executable, "-m", "pip", "check"], capture_output=True, text=True).stdout.strip()
    if out and "No broken requirements" not in out:
        print(f"{WARN} pip check reported:\n{out}\n       (conflicts about transformers/tokenizers are often "
              "harmless with --no-deps; only worry if the model test below fails)")
    else:
        report(OK, "pip check: no broken requirements")

    # 7. Optional end-to-end test
    if args.model:
        try:
            import torch
            from unsloth import FastModel
            from datasets import Dataset
            from trl import SFTConfig, SFTTrainer

            print(f"\nLoading {args.model_name} ...")
            model, tok = FastModel.from_pretrained(
                model_name=args.model_name, dtype=None, max_seq_length=512,
                load_in_4bit=False, full_finetuning=False)
            report(OK, "model + tokenizer loaded")

            model = FastModel.get_peft_model(
                model, finetune_vision_layers=False, finetune_language_layers=True,
                finetune_attention_modules=True, finetune_mlp_modules=True,
                r=8, lora_alpha=8, lora_dropout=0, bias="none", use_rslora=True, random_state=3407)
            report(OK, "LoRA attached (use_rslora accepted)")

            sin = "ශ්‍රී ලංකාව ඉන්දියන් සාගරයේ පිහිටි ද්වීපයකි."
            ids = tok(sin, add_special_tokens=False)["input_ids"]
            print(f"       Sinhala sample -> {len(ids)} tokens for {len(sin)} chars")

            ds = Dataset.from_dict({"text": [(sin + "\n") * 20] * 16})
            cfg = SFTConfig(output_dir="/tmp/envcheck", dataset_text_field="text", max_length=512,
                            packing=True, per_device_train_batch_size=2, gradient_accumulation_steps=1,
                            max_steps=2, learning_rate=5e-5, bf16=True, optim="adamw_8bit",
                            logging_steps=1, save_strategy="no", report_to="none")
            tr = SFTTrainer(model=model, tokenizer=tok, train_dataset=ds, args=cfg)
            tr.train()
            report(OK, f"2 training steps ran; peak VRAM {torch.cuda.max_memory_reserved() / 1024**3:.1f} GB")
        except Exception as e:  # noqa: BLE001
            import traceback
            traceback.print_exc()
            report(BAD, f"model test failed: {type(e).__name__}: {e}")

    print("\n" + ("ALL CHECKS PASSED" if not failures else f"{len(failures)} PROBLEM(S):"))
    for f in failures:
        print("  -", f)
    sys.exit(1 if failures else 0)


if __name__ == "__main__":
    main()