Skip to content

Commit

Permalink
Implemement smoke_test.py for pip wheel jobs
Browse files Browse the repository at this point in the history
This script is run by CI after building the executorch wheel. Before
running this, the job will install the matching torch package as well as
the newly-built executorch package and its dependencies.

For now we test the export of a simple model, and try executing it using
the runtime pybindings.

Test Plan:
```
./install_requirements.sh
python build/packaging/smoke_test.py
```
  • Loading branch information
dbort committed Apr 22, 2024
1 parent aa51190 commit 358a03f
Showing 1 changed file with 91 additions and 5 deletions.
96 changes: 91 additions & 5 deletions build/packaging/smoke_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,13 +5,99 @@
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

"""
This script is run by CI after building the executorch wheel. Before running
this, the job will install the matching torch package as well as the newly-built
executorch package and its dependencies.
"""

# Import this first. If it can't find the torch.so libraries, the dynamic load
# will fail and the process will exit.
from executorch.extension.pybindings import portable_lib # usort: skip

# Import this after importing the ExecuTorch pybindings. If the pybindings
# links against a different torch.so than this uses, there will be a set of
# symbol comflicts; the process will either exit now, or there will be issues
# later in the smoke test.
import torch # usort: skip

# Import everything else later to help isolate the critical imports above.
import os
import tempfile
from typing import Tuple

from executorch.exir import to_edge
from torch.export import export


class LinearModel(torch.nn.Module):
"""Runs Linear on its input, which should have shape [4]."""

def __init__(self):
super().__init__()
self.linear = torch.nn.Linear(4, 2)

def forward(self, x: torch.Tensor):
"""Expects a single tensor of shape [4]."""
return self.linear(x)


def linear_model_inputs() -> Tuple[torch.Tensor]:
"""Returns some example inputs compatible with LinearModel."""
# The model takes a single tensor of shape [4] as an input.
return (torch.ones(4),)


def export_linear_model() -> bytes:
"""Exports LinearModel and returns the .pte data."""

# This helps the exporter understand the shapes of tensors used in the model.
# Since our model only takes one input, this is a one-tuple.
example_inputs = linear_model_inputs()

# Export the pytorch model and process for ExecuTorch.
print("Exporting program...")
exported_program = export(LinearModel(), example_inputs)
print("Lowering to edge...")
edge_program = to_edge(exported_program)
print("Creating ExecuTorch program...")
et_program = edge_program.to_executorch()

return et_program.buffer


def main():
"""
Run ExecuTorch binary smoke tests. This is a placeholder for future tests. See
https://github.com/pytorch/test-infra/wiki/Using-Nova-Reusable-Build-Workflows
for more information about Nova binary workflow.
"""
"""Tests the export and execution of a simple model."""

# If the pybindings loaded correctly, we should be able to ask for the set
# of operators.
ops = portable_lib._get_operator_names()
assert len(ops) > 0, "Empty operator list"
print(f"Found {len(ops)} operators; first element '{ops[0]}'")

# Export LinearModel to .pte data.
pte_data: bytes = export_linear_model()

# Try saving to and loading from a file.
with tempfile.TemporaryDirectory() as tempdir:
pte_file = os.path.join(tempdir, "linear.pte")

# Save the .pte data to a file.
with open(pte_file, "wb") as file:
file.write(pte_data)
print(f"ExecuTorch program saved to {pte_file} ({len(pte_data)} bytes).")

# Load the model from disk.
m = portable_lib._load_for_executorch(pte_file)

# Run the model.
outputs = m.forward(linear_model_inputs())

# Should see a single output with shape [2].
assert len(outputs) == 1, f"Unexpected output length {len(outputs)}: {outputs}"
assert outputs[0].shape == (2,), f"Unexpected output size {outputs[0].shape}"

print("PASS")


if __name__ == "__main__":
Expand Down

0 comments on commit 358a03f

Please sign in to comment.