`torch.onnx.export` not supported with default `dynamo` param

[imMoya]

Summary

When exporting a preprocessing pipeline (export_module), as the one described below:

torch.onnx.export(
    export_module,
    (dummy_xyz, dummy_normals, dummy_params),
    output_path,
    input_names=["xyz", "normals", "params"],
    output_names=["xyz_scaled", "normals_scaled", "params_scaled"],
    dynamic_axes={
        "xyz": {0: "batch_size", 1: "num_points"},
        "normals": {0: "batch_size", 1: "num_points"},
        "params": {0: "batch_size"},
        "xyz_scaled": {0: "batch_size", 1: "num_points"},
        "normals_scaled": {0: "batch_size", 1: "num_points"},
        "params_scaled": {0: "batch_size"},
    },
    opset_version=18,
    do_constant_folding=True,
)

The following TorchExportError is raised

TorchExportError: Failed to export the model with torch.export. This is step
1/3 of exporting the model to ONNX. Next steps:
- Modify the model code for `torch.export.export` to succeed. Refer to 
https://pytorch.org/docs/stable/generated/exportdb/index.html for more 
information.
- Debug `torch.export.export` and submit a PR to PyTorch.
- Create an issue in the PyTorch GitHub repository against the 
*torch.export* component and attach the full error stack as well as 
reproduction scripts.

## Exception summary

<class 'KeyError'>: 'getpwuid(): uid not found: 501'

Workaround
To circumvent this, the option dynamo=False must be set to the export function in order to avoid the error.

Why is this needed?

Use default export option from PyTorch (employs TorchDynamo instead of TorchScript)

Usage example

Minimal example of model export

def evaluate(inputs: Any) -> Any:
    import torch
    import torch.nn as nn

    class SimpleMLP(nn.Module):
        def __init__(self, in_features=4, hidden=32, out_features=1):
            super().__init__()
            self.net = nn.Sequential(
                nn.Linear(in_features, hidden),
                nn.ReLU(),
                nn.Linear(hidden, out_features),
            )

        def forward(self, x):
            return self.net(x)

    def export_onnx(
        model: nn.Module,
        onnx_path: str = "simple_mlp.onnx",
        in_features: int = 16,
    ):
        model.eval()

        # Dummy input defines input shape for tracing
        x = torch.randn(1, in_features)

        torch.onnx.export(
            model,
            x,
            onnx_path,
            export_params=True,
            opset_version=17,
            do_constant_folding=True,
            input_names=["input"],
            output_names=["output"],
            dynamic_axes={
                "input": {0: "batch"},
                "output": {0: "batch"},
            },
        )
        return onnx_path

    model = SimpleMLP(in_features=4, hidden=32, out_features=1)
    path = export_onnx(model, "simple_mlp.onnx", in_features=4)
    return path

[dionhaefner]

@xalelax may be time to ensure our OS users actually exist?

We could add an entrypoint.sh to our Docker images that essentially calls useradd tesseractor <myuid>:<mygid>, but the annoying part is that this would require a world-writable /etc/passwd I think? Which may not be too bad but seems a bit iffy.

[dionhaefner]

@imMoya please post a full reproducer next time that we can drop in a folder and run tesseract build on. I've created one for you this time:

exporter.zip

[xalelax]

✅ Ran @dionhaefner 's exporter with current main with dynamo=True and it worked

$ tesseract run exporter:latest apply '{"inputs": {}}'
/tesseract/tesseract_api.py:54: UserWarning: # 'dynamic_axes' is not recommended when dynamo=True, and may lead to 'torch._dynamo.exc.UserError: Constraints violated.' Supply the 'dynamic_shapes' argument instead if export is unsuccessful.
  torch.onnx.export(
W1231 09:57:20.501000 1 torch/onnx/_internal/exporter/_compat.py:114] Setting ONNX exporter to use operator set version 18 because the requested opset_version 17 is a lower version than we have implementations for. Automatic version conversion will be performed, which may not be successful at converting to the requested version. If version conversion is unsuccessful, the opset version of the exported model will be kept at 18. Please consider setting opset_version >=18 to leverage latest ONNX features
W1231 09:57:20.978000 1 torch/onnx/_internal/exporter/_registration.py:107] torchvision is not installed. Skipping torchvision::nms
The model version conversion is not supported by the onnxscript version converter and fallback is enabled. The model will be converted using the onnx C API (target version: 17).
[torch.onnx] Obtain model graph for `SimpleMLP([...]` with `torch.export.export(..., strict=False)`...
[torch.onnx] Obtain model graph for `SimpleMLP([...]` with `torch.export.export(..., strict=False)`... ✅
[torch.onnx] Run decomposition...
[torch.onnx] Run decomposition... ✅
[torch.onnx] Translate the graph into ONNX...
[torch.onnx] Translate the graph into ONNX... ✅
{"path":"simple_mlp.onnx"}