put executorch in the extras_require (#4324)

Author: lanluo-nvidia

.github/scripts/verify-executorch-reference-runner.sh

@@ -207,11 +207,14 @@ if ! "${python_executable}" - <<'PY'
 import importlib
 import importlib.util
 
-missing = [
-    name
-    for name in ("yaml", "torch", "torch_tensorrt", "executorch.exir")
-    if importlib.util.find_spec(name) is None
-]
+missing = []
+for name in ("yaml", "torch", "torch_tensorrt", "executorch.exir"):
+    try:
+        spec = importlib.util.find_spec(name)
+    except ModuleNotFoundError:
+        spec = None
+    if spec is None:
+        missing.append(name)
 if missing:
     raise SystemExit(
         "Missing Python package(s) required to export the .pte and build the runner: "

.github/workflows/executorch-static-linux.yml

@@ -96,5 +96,5 @@ jobs:
         export EXECUTORCH_SOURCE_DIR="$(dirname "${executorch_cmake_location%%:*}")"
         export EXECUTORCH_ROOT="${EXECUTORCH_SOURCE_DIR}"
         # this is to verify the end user's workflow
-        python -m pip install pyyaml
+        python -m pip install pyyaml "executorch>=1.3.1"
         .github/scripts/verify-executorch-reference-runner.sh

py/torch_tensorrt/_compile.py

@@ -767,8 +767,8 @@ def save(
     if output_format == "executorch" and not _has_executorch_exir():
         raise ImportError(
             "Saving in ExecuTorch format requires the executorch package "
-            "with executorch.exir. Install executorch to use "
-            "output_format='executorch'."
+            "with executorch.exir. Install with: pip install "
+            "\"executorch\" to use output_format='executorch'."
         )
 
     def _all_are_input_objects(obj: Any) -> bool:
@@ -1314,8 +1314,8 @@ def _save_as_executorch(exp_program: Any, file_path: str, **kwargs: Any) -> None
         from executorch.exir import to_edge_transform_and_lower
     except ImportError:
         raise ImportError(
-            "ExecuTorch is not installed. Please install it to use output_format='executorch'. "
-            "See https://pytorch.org/executorch/stable/getting-started-setup.html"
+            "ExecuTorch is not installed. Install with: pip install "
+            "\"executorch\" to use output_format='executorch'."
         )
     import torch_tensorrt.dynamo.runtime.meta_ops.register_meta_ops  # noqa: F401
     from torch_tensorrt.executorch import (

py/torch_tensorrt/executorch/__init__.py

@@ -18,7 +18,7 @@ def __getattr__(name: str) -> NoReturn:
         raise ImportError(
             f"Cannot access torch_tensorrt.executorch.{name}: "
             "ExecuTorch with executorch.exir is required. "
-            "Install with: pip install executorch"
+            'Install with: pip install "executorch"'
         )
 
     __all__ = [

setup.py

@@ -99,7 +99,11 @@ def load_dep_info():
 CI_BUILD = False
 USE_TRT_RTX = False
 
-EXECUTORCH_REQUIREMENT = "executorch>=1.2.0"
+EXECUTORCH_REQUIREMENT = "executorch>=1.3.1"
+EXTRAS_REQUIRE = {
+    "executorch": [EXECUTORCH_REQUIREMENT],
+    "all": [EXECUTORCH_REQUIREMENT],
+}
 
 
 if "--use-rtx" in sys.argv:
@@ -608,9 +612,7 @@ def run(self):
     )
 
 package_data = {}
-executorch_header_package_data = (
-    [] if IS_DLFW_CI else ["include/torch_tensorrt/executorch/*.h"]
-)
+executorch_header_package_data = ["include/torch_tensorrt/executorch/*.h"]
 
 if not (PY_ONLY or NO_TS):
     tensorrt_x86_64_external_dir = (
@@ -855,7 +857,6 @@ def get_x86_64_requirements(base_requirements):
             ]
         else:
             requirements = requirements + [
-                EXECUTORCH_REQUIREMENT,
                 "tensorrt>=11.0.0,<11.1.0",
             ]
             cuda_version = torch.version.cuda
@@ -914,6 +915,7 @@ def get_requirements():
     },
     zip_safe=False,
     install_requires=get_requirements(),
+    extras_require=EXTRAS_REQUIRE,
     packages=packages,
     package_dir=package_dir,
     include_package_data=False,

tests/py/dynamo/executorch/test_api.py

@@ -1,5 +1,7 @@
+import ast
 import importlib
 import sys
+from pathlib import Path
 
 import pytest
 import torch
@@ -18,7 +20,7 @@ def fake_find_spec(name, package=None):
     monkeypatch.setattr(importlib.util, "find_spec", fake_find_spec)
     module = importlib.import_module("torch_tensorrt.executorch")
 
-    with pytest.raises(ImportError, match="ExecuTorch.*required"):
+    with pytest.raises(ImportError, match=r"torch_tensorrt\[executorch\]"):
         _ = module.TensorRTBackend
 
     sys.modules.pop("torch_tensorrt.executorch", None)
@@ -39,7 +41,7 @@ def fake_find_spec(name, package=None):
 
     from torch_tensorrt._compile import save
 
-    with pytest.raises(ImportError, match="Saving in ExecuTorch format requires"):
+    with pytest.raises(ImportError, match=r"torch_tensorrt\[executorch\]"):
         save(
             torch.nn.Linear(1, 1),
             str(tmp_path / "model.pte"),
@@ -53,3 +55,92 @@ def test_public_api_symbols_present():
     assert "get_edge_compile_config" in module.__all__
     assert "TensorRTPartitioner" in module.__all__
     assert "TensorRTBackend" in module.__all__
+
+
+_REPO_ROOT = Path(__file__).resolve().parents[4]
+_SETUP_PY = _REPO_ROOT / "setup.py"
+
+
+def _setup_tree():
+    return ast.parse(_SETUP_PY.read_text(encoding="utf-8"))
+
+
+def _assignment_value(tree, name):
+    for node in tree.body:
+        if isinstance(node, ast.Assign) and any(
+            isinstance(target, ast.Name) and target.id == name
+            for target in node.targets
+        ):
+            return node.value
+    raise AssertionError(f"Could not find assignment for {name}")
+
+
+def _function_def(tree, name):
+    for node in tree.body:
+        if isinstance(node, ast.FunctionDef) and node.name == name:
+            return node
+    raise AssertionError(f"Could not find function {name}")
+
+
+@pytest.mark.unit
+def test_packaging_declares_executorch_extra():
+    tree = _setup_tree()
+    extras = _assignment_value(tree, "EXTRAS_REQUIRE")
+    assert isinstance(extras, ast.Dict)
+
+    extras_by_name = {
+        key.value: value
+        for key, value in zip(extras.keys, extras.values)
+        if isinstance(key, ast.Constant)
+    }
+    for extra_name in ("executorch", "all"):
+        assert extra_name in extras_by_name
+        requirements = extras_by_name[extra_name]
+        assert isinstance(requirements, ast.List)
+        assert any(
+            isinstance(requirement, ast.Name)
+            and requirement.id == "EXECUTORCH_REQUIREMENT"
+            for requirement in requirements.elts
+        )
+
+    setup_call = next(
+        node
+        for node in ast.walk(tree)
+        if isinstance(node, ast.Call)
+        and isinstance(node.func, ast.Name)
+        and node.func.id == "setup"
+    )
+    extras_keyword = next(
+        (keyword for keyword in setup_call.keywords if keyword.arg == "extras_require"),
+        None,
+    )
+    assert extras_keyword is not None
+    assert isinstance(extras_keyword.value, ast.Name)
+    assert extras_keyword.value.id == "EXTRAS_REQUIRE"
+
+
+@pytest.mark.unit
+def test_executorch_is_not_base_install_requirement():
+    tree = _setup_tree()
+    for function_name in (
+        "get_jetpack_requirements",
+        "get_sbsa_requirements",
+        "get_x86_64_requirements",
+        "get_requirements",
+    ):
+        function = _function_def(tree, function_name)
+        assert not any(
+            isinstance(node, ast.Name) and node.id == "EXECUTORCH_REQUIREMENT"
+            for node in ast.walk(function)
+        )
+
+
+@pytest.mark.unit
+def test_executorch_headers_are_not_dlfw_gated():
+    tree = _setup_tree()
+    header_package_data = _assignment_value(tree, "executorch_header_package_data")
+    assert isinstance(header_package_data, ast.List)
+    assert not any(
+        isinstance(node, ast.Name) and node.id == "IS_DLFW_CI"
+        for node in ast.walk(header_package_data)
+    )