feat: Validate endpoint argument names before building (#95)

BREAKING CHANGE: Custom argument names for any endpoint are no longer
supported. All endpoint argument names and order must exactly match
standard.

#### Relevant issue or PR
- #84

#### Description of changes
We already validate whether endpoint signature have the correct number
of arguments in `sdk.api_parse`. In addition to suppressing unnecessary
traceback and improving the error message when this is violated, this PR
provides validation of the endpoint signature argument _names_ and order
before building. This guards against the possibility that arguments may
have been misordered leading to potential silent fails and guarantees
that natural dev assumptions on the endpoint signatures hold.

Note it is not possible to provide this with `tesseract-runtime check`
(or other endpoints) as the runtime does not have access to the sdk. I'm
comfortable with this as we can add runtime errors in a different way by
reviving #80. Alternatively, we move `api_parse` to `runtime` allowing
it to be accessed by `tesseract-runtime check`.

#### Missing argument error message
```
 [-] Error validating tesseract_api.py: jacobian_vector_product must have 4 arguments: inputs, jvp_inputs, jvp_outputs, tangent_vector.
However, tesseract_api.py specifies 3 arguments: inputs, jvp_outputs, tangent_vector.
 [x] Aborting
```

#### Wrong argument name error message
```
 [-] Error validating tesseract_api.py: The second argument of jacobian_vector_product must be named jvp_inputs, but tesseract_api.py has named it jvp_input.
 [x] Aborting
```

#### Testing done
- [x] Removed an input, error message as above
- [x] Misspelled an input, error message as above 
- [x] CI passes

#### License

- [x] By submitting this pull request, I confirm that my contribution is
made under the terms of the [Apache 2.0
license](https://pasteurlabs.github.io/tesseract/LICENSE).
- [x] I sign the Developer Certificate of Origin below by adding my name
and email address to the `Signed-off-by` line.

<details>
<summary><b>Developer Certificate of Origin</b></summary>

```text
Developer Certificate of Origin
Version 1.1

Copyright (C) 2004, 2006 The Linux Foundation and its contributors.

Everyone is permitted to copy and distribute verbatim copies of this
license document, but changing it is not allowed.


Developer's Certificate of Origin 1.1

By making a contribution to this project, I certify that:

(a) The contribution was created in whole or in part by me and I
    have the right to submit it under the open source license
    indicated in the file; or

(b) The contribution is based upon previous work that, to the best
    of my knowledge, is covered under an appropriate open source
    license and I have the right under that license to submit that
    work with modifications, whether created in whole or in part
    by me, under the same open source license (unless I am
    permitted to submit under a different license), as indicated
    in the file; or

(c) The contribution was provided directly to me by some other
    person who certified (a), (b) or (c) and I have not modified
    it.

(d) I understand and agree that this project and the contribution
    are public and that a record of the contribution (including all
    personal information I submit with it, including my sign-off) is
    maintained indefinitely and may be redistributed consistent with
    this project or the open source license(s) involved.
```

</details>

Signed-off-by: Jonathan Brodrick <jonathan.brodrick@simulation.science>

Author: jpbrodrick89

examples/vectoradd_torch/tesseract_api.py

@@ -150,14 +150,14 @@ def jacobian_vector_product(
     inputs: InputSchema,
     jvp_inputs: set[str],
     jvp_outputs: set[str],
-    tangent: dict[str, Any],
+    tangent_vector: dict[str, Any],
 ):
-    # Make ordering of tangent identical to jvp_inputs
-    tangent = {key: tangent[key] for key in jvp_inputs}
+    # Make ordering of tangent_vector identical to jvp_inputs
+    tangent_vector = {key: tangent_vector[key] for key in jvp_inputs}
 
     # convert all numbers and arrays to torch tensors
     tensor_inputs = tree_map(to_tensor, inputs.model_dump())
-    pos_tangent = tree_map(to_tensor, tangent).values()
+    pos_tangent = tree_map(to_tensor, tangent_vector).values()
 
     # flatten the dictionaries such that they can be accessed by paths
     pos_inputs = flatten_with_paths(tensor_inputs, jvp_inputs).values()

tesseract_core/sdk/api_parse.py

@@ -21,17 +21,44 @@ class _ApiObject(NamedTuple):
     name: str
     expected_type: type
     num_args: int | None = None
+    arg_names: tuple[str, ...] | None = None
     optional: bool = False
 
 
+ORDINALS = ["first", "second", "third", "fourth", "fifth", "sixth", "seventh", "eighth"]
+
 EXPECTED_OBJECTS = (
-    _ApiObject("apply", ast.FunctionDef, 1),
+    _ApiObject("apply", ast.FunctionDef, 1, arg_names=("inputs",)),
     _ApiObject("InputSchema", ast.ClassDef),
     _ApiObject("OutputSchema", ast.ClassDef),
-    _ApiObject("jacobian", ast.FunctionDef, 3, optional=True),
-    _ApiObject("jacobian_vector_product", ast.FunctionDef, 4, optional=True),
-    _ApiObject("vector_jacobian_product", ast.FunctionDef, 4, optional=True),
-    _ApiObject("abstract_eval", ast.FunctionDef, 1, optional=True),
+    _ApiObject(
+        "jacobian",
+        ast.FunctionDef,
+        3,
+        arg_names=("inputs", "jac_inputs", "jac_outputs"),
+        optional=True,
+    ),
+    _ApiObject(
+        "jacobian_vector_product",
+        ast.FunctionDef,
+        4,
+        arg_names=("inputs", "jvp_inputs", "jvp_outputs", "tangent_vector"),
+        optional=True,
+    ),
+    _ApiObject(
+        "vector_jacobian_product",
+        ast.FunctionDef,
+        4,
+        arg_names=("inputs", "vjp_inputs", "vjp_outputs", "cotangent_vector"),
+        optional=True,
+    ),
+    _ApiObject(
+        "abstract_eval",
+        ast.FunctionDef,
+        1,
+        arg_names=("abstract_inputs",),
+        optional=True,
+    ),
 )
 
 
@@ -112,8 +139,8 @@ class ValidationError(Exception):
     pass
 
 
-def _get_func_argnums(func: ast.FunctionDef) -> int:
-    """Get the number of the arguments of a function node.
+def _get_func_argnames(func: ast.FunctionDef) -> tuple[str, ...]:
+    """Get the names of the arguments of a function node.
 
     See:
     https://docs.python.org/3/library/ast.html#ast.FunctionDef
@@ -128,7 +155,7 @@ def _get_func_argnums(func: ast.FunctionDef) -> int:
         raise ValidationError(
             f"Function {func.name} must not have positional-only arguments"
         )
-    return len(func_args.args)
+    return tuple(arg.arg for arg in func_args.args)
 
 
 def validate_tesseract_api(src_dir: Path) -> None:
@@ -190,8 +217,23 @@ def validate_tesseract_api(src_dir: Path) -> None:
             )
 
         if obj.num_args is not None:
-            if _get_func_argnums(toplevel_objects[obj.name]) != obj.num_args:
-                raise ValidationError(f"{obj.name} must have {obj.num_args} arguments")
+            func_argnames = _get_func_argnames(toplevel_objects[obj.name])
+            func_argnums = len(func_argnames)
+            if func_argnums != obj.num_args:
+                raise ValidationError(
+                    f"{obj.name} must have {obj.num_args} arguments: {', '.join(obj.arg_names)}.\n"
+                    f"However, {tesseract_api_location} specifies {func_argnums} "
+                    f"arguments: {', '.join(func_argnames)}."
+                )
+            msgs = []
+            for i in range(obj.num_args):
+                if func_argnames[i] != obj.arg_names[i]:
+                    msgs.append(
+                        f"The {ORDINALS[i]} argument (argument {i}) of {obj.name} must be named {obj.arg_names[i]}, "
+                        f"but {tesseract_api_location} has named it {func_argnames[i]}."
+                    )
+            if msgs:
+                raise ValidationError("\n".join(msgs))
 
     # Check InputSchema and OutputSchema are pydantic BaseModels
     for schema in ("InputSchema", "OutputSchema"):

tesseract_core/sdk/cli.py

@@ -33,6 +33,7 @@
     EXPECTED_OBJECTS,
     TesseractBuildConfig,
     TesseractConfig,
+    ValidationError,
     get_non_base_fields_in_tesseract_config,
 )
 from .exceptions import UserError
@@ -307,6 +308,8 @@ def build_image(
         raise UserError(f"Input error building Tesseract: {e}") from e
     except PermissionError as e:
         raise UserError(f"Permission denied: {e}") from e
+    except ValidationError as e:
+        raise UserError(f"Error validating tesseract_api.py: {e}") from e
 
     if generate_only:
         # output is the path to the build context

tesseract_core/sdk/templates/pytorch/tesseract_api.py

@@ -109,14 +109,14 @@ def jacobian_vector_product(
     inputs: InputSchema,
     jvp_inputs: set[str],
     jvp_outputs: set[str],
-    tangent: dict[str, Any],
+    tangent_vector: dict[str, Any],
 ):
-    # Make ordering of tangent identical to jvp_inputs
-    tangent = {key: tangent[key] for key in jvp_inputs}
+    # Make ordering of tangent_vector identical to jvp_inputs
+    tangent_vector = {key: tangent_vector[key] for key in jvp_inputs}
 
     # convert all numbers and arrays to torch tensors
     tensor_inputs = tree_map(to_tensor, inputs.model_dump())
-    pos_tangent = tree_map(to_tensor, tangent).values()
+    pos_tangent = tree_map(to_tensor, tangent_vector).values()
 
     # flatten the dictionaries such that they can be accessed by paths
     pos_inputs = flatten_with_paths(tensor_inputs, jvp_inputs).values()

tests/runtime_tests/test_core.py

@@ -105,7 +105,7 @@ def vector_jacobian_product(
 
             return out
 
-        def abstract_eval(self, inputs):
+        def abstract_eval(self, abstract_inputs):
             return {
                 "result_seq": [{"shape": (10,), "dtype": "float32"}] * 3,
                 "result_arr": {"shape": (5,), "dtype": "float32"},

tests/runtime_tests/test_serve.py

@@ -179,7 +179,7 @@ def apply(input: InputSchema) -> OutputSchema:
         assert threading.current_thread() == threading.main_thread()
         return OutputSchema()
 
-    def abstract_eval(inputs: dict) -> dict:
+    def abstract_eval(abstract_inputs: dict) -> dict:
         pass
     """
     )