Skip to content

Commit a88b7d1

Browse files
authored
Metal backend: add export-time code signing for Hardened Runtime (#18768)
The Metal AOTI backend extracts a compiled .so from the .pte at runtime and dlopen's it. macOS Hardened Runtime rejects unsigned dlopen'd code, making Metal-backend .pte files unusable in notarized apps. Add a `codesign_so` hook to `AotiBackend` (no-op by default) that runs after AOTInductor compilation and before the .so is packed into the .pte. `MetalBackend` overrides it to run `codesign` when a `codesign_identity` compile spec is provided. Wire `--codesign-identity` through the Voxtral Realtime and Parakeet export scripts.
1 parent 196a4aa commit a88b7d1

5 files changed

Lines changed: 99 additions & 11 deletions

File tree

backends/aoti/aoti_backend.py

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,19 @@ def get_extra_aoti_compile_context_manager(cls):
9292
"""Return extra context manager to apply during aoti_compile stage. By default returns an empty context manager."""
9393
return contextlib.nullcontext()
9494

95+
@classmethod
96+
def codesign_so(cls, so_path: str, compile_specs: List[CompileSpec]) -> None:
97+
"""Sign the compiled .so before packing into .pte.
98+
99+
Called after AOTInductor compilation, before the .so bytes are read
100+
and packed into the named data store. Override in platform-specific
101+
backends to apply code signing (e.g., macOS codesign for Hardened
102+
Runtime compatibility).
103+
104+
Default: no-op.
105+
"""
106+
return
107+
95108
@classmethod
96109
@contextlib.contextmanager
97110
def collect_unsupported_fallback_kernels(cls, missing_fallback_kernels: Set[str]):
@@ -226,6 +239,9 @@ def preprocess(
226239
f"Could not find required files in compiled paths, got {paths}"
227240
)
228241

242+
# Sign the .so for platform-specific requirements (e.g., macOS Hardened Runtime)
243+
cls.codesign_so(so_path, compile_specs)
244+
229245
# Read SO file
230246
with open(so_path, "rb") as f:
231247
so_data = f.read()

backends/apple/metal/metal_backend.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
# This source code is licensed under the BSD-style license found in the
55
# LICENSE file in the root directory of this source tree.
66

7+
import subprocess
78
import typing
89
from typing import Any, Dict, final, List
910

@@ -102,3 +103,16 @@ def get_aoti_compile_options(
102103
inductor_configs["aot_inductor.custom_ops_to_c_shims"] = custom_c_shims
103104

104105
return inductor_configs
106+
107+
@classmethod
108+
def codesign_so(cls, so_path: str, compile_specs: List[CompileSpec]) -> None:
109+
"""Sign the compiled .so for macOS Hardened Runtime compatibility.
110+
111+
Only signs if a ``codesign_identity`` compile spec is provided.
112+
Pass ``"-"`` for ad-hoc signing or a Developer ID for distribution.
113+
"""
114+
for spec in compile_specs:
115+
if spec.key == "codesign_identity":
116+
identity = spec.value.decode("utf-8")
117+
subprocess.run(["codesign", "-f", "-s", identity, so_path], check=True)
118+
return

backends/apple/metal/tests/test_modules.py

Lines changed: 28 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
from executorch.backends.apple.metal.metal_backend import MetalBackend
2828
from executorch.backends.apple.metal.metal_partitioner import MetalPartitioner
2929
from executorch.exir import to_edge_transform_and_lower
30+
from executorch.exir.backend.compile_spec_schema import CompileSpec
3031
from torch import nn
3132
from torch.export import export
3233
from torch.nn.attention import SDPBackend
@@ -959,21 +960,25 @@ def linear_filter(m, fqn):
959960

960961

961962
def export_model_to_metal(
962-
model: nn.Module, example_inputs: Tuple[torch.Tensor, ...]
963+
model: nn.Module,
964+
example_inputs: Tuple[torch.Tensor, ...],
965+
codesign_identity: Optional[str] = None,
963966
) -> Any:
964967
"""Export model through the Metal backend pipeline."""
965968
method_name = "forward"
966969

970+
compile_specs = [MetalBackend.generate_method_name_compile_spec(method_name)]
971+
if codesign_identity:
972+
compile_specs.append(
973+
CompileSpec("codesign_identity", codesign_identity.encode("utf-8"))
974+
)
975+
967976
with torch.nn.attention.sdpa_kernel([SDPBackend.MATH]), torch.no_grad():
968977
aten_dialect = export(model, example_inputs, strict=False)
969978

970979
edge_program = to_edge_transform_and_lower(
971980
aten_dialect,
972-
partitioner=[
973-
MetalPartitioner(
974-
[MetalBackend.generate_method_name_compile_spec(method_name)]
975-
)
976-
],
981+
partitioner=[MetalPartitioner(compile_specs)],
977982
)
978983

979984
executorch_program = edge_program.to_executorch()
@@ -1326,6 +1331,23 @@ def run_test_in_directory(test_dir: Path) -> None:
13261331
with tempfile.TemporaryDirectory() as tmpdir:
13271332
run_test_in_directory(Path(tmpdir))
13281333

1334+
@unittest.skipIf(SKIP_EXPORT_TESTS, SKIP_REASON)
1335+
@unittest.skipIf(not IS_MACOS, "codesign requires macOS")
1336+
def test_codesign_export(self):
1337+
"""Test that export with codesign_identity='-' signs the .so and succeeds."""
1338+
model, example_inputs = get_model_and_inputs("add", dtype=torch.float32)
1339+
# codesign -f -s - runs during export; check=True means CalledProcessError
1340+
# is raised (and the test fails) if signing fails.
1341+
program = export_model_to_metal(model, example_inputs, codesign_identity="-")
1342+
self.assertGreater(len(program.buffer), 0)
1343+
1344+
@unittest.skipIf(SKIP_EXPORT_TESTS, SKIP_REASON)
1345+
def test_export_without_codesign(self):
1346+
"""Test that export without codesign_identity skips signing."""
1347+
model, example_inputs = get_model_and_inputs("add", dtype=torch.float32)
1348+
program = export_model_to_metal(model, example_inputs)
1349+
self.assertGreater(len(program.buffer), 0)
1350+
13291351

13301352
# =============================================================================
13311353
# Dynamically generate test methods for each module and dtype in MODULE_REGISTRY

examples/models/parakeet/export_parakeet_tdt.py

Lines changed: 22 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -498,10 +498,11 @@ def _linear_bias_decomposition(input, weight, bias=None):
498498
return out
499499

500500

501-
def _create_metal_partitioners(programs):
501+
def _create_metal_partitioners(programs, codesign_identity=None):
502502
"""Create Metal partitioners for all programs except preprocessor."""
503503
from executorch.backends.apple.metal.metal_backend import MetalBackend
504504
from executorch.backends.apple.metal.metal_partitioner import MetalPartitioner
505+
from executorch.exir.backend.compile_spec_schema import CompileSpec
505506

506507
print("\nLowering to ExecuTorch with Metal...")
507508

@@ -521,6 +522,10 @@ def _create_metal_partitioners(programs):
521522
partitioner[key] = []
522523
else:
523524
compile_specs = [MetalBackend.generate_method_name_compile_spec(key)]
525+
if codesign_identity:
526+
compile_specs.append(
527+
CompileSpec("codesign_identity", codesign_identity.encode("utf-8"))
528+
)
524529
partitioner[key] = [MetalPartitioner(compile_specs)]
525530
return partitioner, updated_programs
526531

@@ -586,12 +591,18 @@ def _create_vulkan_partitioners(programs, vulkan_force_fp16=False):
586591

587592

588593
def lower_to_executorch(
589-
programs, metadata=None, backend="portable", vulkan_force_fp16=False
594+
programs,
595+
metadata=None,
596+
backend="portable",
597+
vulkan_force_fp16=False,
598+
codesign_identity=None,
590599
):
591600
if backend == "xnnpack":
592601
partitioner, programs = _create_xnnpack_partitioners(programs)
593602
elif backend == "metal":
594-
partitioner, programs = _create_metal_partitioners(programs)
603+
partitioner, programs = _create_metal_partitioners(
604+
programs, codesign_identity=codesign_identity
605+
)
595606
elif backend == "mlx":
596607
partitioner, programs = _create_mlx_partitioners(programs)
597608
elif backend in ("cuda", "cuda-windows"):
@@ -714,6 +725,13 @@ def main():
714725
)
715726

716727
parser.add_argument("--vulkan_force_fp16", action="store_true")
728+
parser.add_argument(
729+
"--codesign-identity",
730+
default=None,
731+
help="macOS code signing identity for the Metal backend .so. "
732+
"Use '-' for ad-hoc or a Developer ID for notarized apps. "
733+
"If omitted, the .so is not signed.",
734+
)
717735

718736
args = parser.parse_args()
719737

@@ -767,6 +785,7 @@ def main():
767785
metadata=metadata,
768786
backend=args.backend,
769787
vulkan_force_fp16=args.vulkan_force_fp16,
788+
codesign_identity=args.codesign_identity,
770789
)
771790

772791
pte_path = os.path.join(args.output_dir, "model.pte")

examples/models/voxtral_realtime/export_voxtral_rt.py

Lines changed: 19 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -379,7 +379,7 @@ def _linear_bias_decomposition(input, weight, bias=None):
379379
return out
380380

381381

382-
def lower_to_executorch(programs, metadata, backend="xnnpack"):
382+
def lower_to_executorch(programs, metadata, backend="xnnpack", codesign_identity=None):
383383
"""Lower exported programs to ExecuTorch."""
384384
transform_passes = None
385385

@@ -397,6 +397,7 @@ def lower_to_executorch(programs, metadata, backend="xnnpack"):
397397
elif backend == "metal":
398398
from executorch.backends.apple.metal.metal_backend import MetalBackend
399399
from executorch.backends.apple.metal.metal_partitioner import MetalPartitioner
400+
from executorch.exir.backend.compile_spec_schema import CompileSpec
400401

401402
print("\nLowering to ExecuTorch with Metal...")
402403

@@ -411,6 +412,10 @@ def lower_to_executorch(programs, metadata, backend="xnnpack"):
411412
partitioner = {}
412413
for key in programs:
413414
compile_specs = [MetalBackend.generate_method_name_compile_spec(key)]
415+
if codesign_identity:
416+
compile_specs.append(
417+
CompileSpec("codesign_identity", codesign_identity.encode("utf-8"))
418+
)
414419
partitioner[key] = [MetalPartitioner(compile_specs)]
415420
elif backend in ("cuda", "cuda-windows"):
416421
from executorch.backends.cuda.cuda_backend import CudaBackend
@@ -577,6 +582,13 @@ def main():
577582
choices=["fp32", "bf16"],
578583
help="Model dtype (default: fp32).",
579584
)
585+
parser.add_argument(
586+
"--codesign-identity",
587+
default=None,
588+
help="macOS code signing identity for the Metal backend .so. "
589+
"Use '-' for ad-hoc or a Developer ID for notarized apps. "
590+
"If omitted, the .so is not signed.",
591+
)
580592
args = parser.parse_args()
581593
backend_for_export = "cuda" if args.backend == "cuda-windows" else args.backend
582594

@@ -638,7 +650,12 @@ def main():
638650
metadata["delay_tokens"] = args.delay_tokens
639651

640652
# Lower
641-
et = lower_to_executorch(programs, metadata, backend=args.backend)
653+
et = lower_to_executorch(
654+
programs,
655+
metadata,
656+
backend=args.backend,
657+
codesign_identity=args.codesign_identity,
658+
)
642659

643660
# Save
644661
pte_path = os.path.join(args.output_dir, "model.pte")

0 commit comments

Comments
 (0)