File size: 5,518 Bytes
2adb41a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
"""
TrackChain CLI Step Skip-Logic Evaluator (tc.v1 SOTA).
Used by ml/scripts/run.sh to evaluate manifest fingerprints and quality gates.

Usage:
  python ml/scripts/should_skip.py --step train_yolo [--force]
  Exit code 0: Skip the step (valid checkpoint & manifest found)
  Exit code 1: Run the step (missing, modified config/dataset, or failed gate)
"""

import sys
import argparse
from pathlib import Path

# Add repo root to python path
repo_root = Path(__file__).resolve().parent.parent.parent
if str(repo_root) not in sys.path:
    sys.path.insert(0, str(repo_root))

from ml.core.manifest import (
    compute_file_sha256,
    compute_config_hash,
    compute_dataset_hash,
    check_training_skip,
    check_export_skip,
    load_manifest,
)


def evaluate_step(step_name: str, force: bool = False) -> bool:
    """Returns True if step should be SKIPPED, False if step should RUN."""
    if force:
        print(f"[{step_name}] Force flag enabled. Running step.")
        return False

    if step_name == "train_yolo":
        # Resolve dataset yaml
        data_yaml = repo_root / "data" / "external" / "rail_defects_expanded" / "data.yaml"
        if not data_yaml.exists():
            data_yaml = repo_root / "data" / "external" / "rail_defects" / "data.yaml"

        config_yaml = repo_root / "ml" / "configs" / "detector.yaml"
        ckpt = repo_root / "artifacts" / "checkpoints" / "vision" / "yolov8n_rail_best.pt"
        manifest_p = repo_root / "artifacts" / "checkpoints" / "vision" / "yolo_manifest.json"

        cfg_hash = compute_config_hash(config_yaml, overrides={"imgsz": 960, "freeze": 10, "dropout": 0.1, "close_mosaic": 10})
        data_hash = compute_dataset_hash(data_yaml)

        skip, reason = check_training_skip(
            manifest_path=manifest_p,
            checkpoint_path=ckpt,
            current_config_hash=cfg_hash,
            current_dataset_hash=data_hash,
            min_metrics={"mAP50": 0.25},
            force=force,
        )
        if skip:
            print(f"[{step_name}] SKIP: {reason}")
            return True
        else:
            print(f"[{step_name}] RUN: {reason}")
            return False

    elif step_name == "export_yolo_onnx":
        src_pt = repo_root / "artifacts" / "checkpoints" / "vision" / "yolov8n_rail_best.pt"
        export_file = repo_root / "artifacts" / "exports" / "yolov8n_rail_best.onnx"
        manifest_p = repo_root / "artifacts" / "exports" / "export_manifest.json"

        skip, reason = check_export_skip(
            manifest_path=manifest_p,
            export_path=export_file,
            source_model_path=src_pt,
            export_format="onnx",
            force=force,
        )
        if skip:
            print(f"[{step_name}] SKIP: {reason}")
            return True
        else:
            print(f"[{step_name}] RUN: {reason}")
            return False

    elif step_name == "export_yolo_int8":
        src_pt = repo_root / "artifacts" / "checkpoints" / "vision" / "yolov8n_rail_best.pt"
        export_file = repo_root / "artifacts" / "exports" / "yolov8n_rail_best_int8.onnx"
        manifest_p = repo_root / "artifacts" / "exports" / "export_manifest.json"

        skip, reason = check_export_skip(
            manifest_path=manifest_p,
            export_path=export_file,
            source_model_path=src_pt,
            export_format="int8",
            force=force,
        )
        if skip:
            print(f"[{step_name}] SKIP: {reason}")
            return True
        else:
            print(f"[{step_name}] RUN: {reason}")
            return False

    elif step_name == "train_patchcore":
        ckpt = repo_root / "artifacts" / "checkpoints" / "vision" / "patchcore_memory_bank.npz"
        calib = repo_root / "artifacts" / "calibration" / "patchcore_calibration.json"
        if ckpt.exists() and calib.exists():
            print(f"[{step_name}] SKIP: PatchCore memory bank and calibration exist.")
            return True
        print(f"[{step_name}] RUN: Memory bank or calibration missing.")
        return False

    elif step_name == "train_bilstm":
        ckpt = repo_root / "artifacts" / "checkpoints" / "geometry" / "bilstm_fault_typing_enhanced.pt"
        calib = repo_root / "artifacts" / "calibration" / "bilstm_temp.json"
        if ckpt.exists() and calib.exists():
            print(f"[{step_name}] SKIP: Bi-LSTM checkpoint and calibration exist.")
            return True
        print(f"[{step_name}] RUN: Bi-LSTM weights or calibration missing.")
        return False

    elif step_name == "train_vae":
        ckpt = repo_root / "artifacts" / "checkpoints" / "geometry" / "sequence_vae_enhanced.pt"
        calib = repo_root / "artifacts" / "calibration" / "vae_calibration.json"
        if ckpt.exists() and calib.exists():
            print(f"[{step_name}] SKIP: Seq-VAE checkpoint and EVT calibration exist.")
            return True
        print(f"[{step_name}] RUN: Seq-VAE weights or calibration missing.")
        return False

    # Default: Run step
    return False


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Evaluate step skip-logic using manifests.")
    parser.add_argument("--step", required=True, help="Step name to evaluate")
    parser.add_argument("--force", action="store_true", help="Force execution")
    args = parser.parse_args()

    should_skip = evaluate_step(step_name=args.step, force=args.force)
    # Exit 0 if skipping, Exit 1 if running
    sys.exit(0 if should_skip else 1)