feat: Add run_id parameter to Tesseract.apply and friends (#352)

<!--
Please use a PR title that conforms to *conventional commits*:
"<commit_type>: Describe your change"; for example: "fix: prevent race
condition". Some other commit types are: fix, feat, ci, doc, refactor...
For a full list of commit types visit
https://www.conventionalcommits.org/en/v1.0.0/
-->

#### Relevant issue or PR
<!-- If the changes resolve an issue or follow some other PR, link to
them here. Only link something if it is directly relevant. -->

#### Description of changes
Adds an optional run_id parameter to `Tesseract.apply`, jacobian, vjp
and jvp methods. If passed, this can be used to correlate the call with
the outputs under `output_path/run_{run_id}`.

#### Testing done
<!-- Describe how the changes were tested; e.g., "CI passes", "Tested
manually in stagingrepo#123", screenshots of a terminal session that
verify the changes, or any other evidence of testing the changes. -->

Author: johnbcoughlin

tesseract_core/sdk/tesseract.py

@@ -310,17 +310,19 @@ def available_endpoints(self) -> list[str]:
         return [endpoint.lstrip("/") for endpoint in self.openapi_schema["paths"]]
 
     @requires_client
-    def apply(self, inputs: dict) -> dict:
+    def apply(self, inputs: dict, run_id: str | None = None) -> dict:
         """Run apply endpoint.
 
         Args:
             inputs: a dictionary with the inputs.
+            run_id: a string to identify the run. Run outputs will be located
+                    in a directory suffixed with this id.
 
         Returns:
             dictionary with the results.
         """
         payload = {"inputs": inputs}
-        return self._client.run_tesseract("apply", payload)
+        return self._client.run_tesseract("apply", payload, run_id)
 
     @requires_client
     def abstract_eval(self, abstract_inputs: dict) -> dict:
@@ -346,14 +348,20 @@ def health(self) -> dict:
 
     @requires_client
     def jacobian(
-        self, inputs: dict, jac_inputs: list[str], jac_outputs: list[str]
+        self,
+        inputs: dict,
+        jac_inputs: list[str],
+        jac_outputs: list[str],
+        run_id: str | None = None,
     ) -> dict:
         """Calculate the Jacobian of (some of the) outputs w.r.t. (some of the) inputs.
 
         Args:
             inputs: a dictionary with the inputs.
             jac_inputs: Inputs with respect to which derivatives will be calculated.
             jac_outputs: Outputs which will be differentiated.
+            run_id: a string to identify the run. Run outputs will be located
+                    in a directory suffixed with this id.
 
         Returns:
             dictionary with the results.
@@ -366,7 +374,7 @@ def jacobian(
             "jac_inputs": jac_inputs,
             "jac_outputs": jac_outputs,
         }
-        return self._client.run_tesseract("jacobian", payload)
+        return self._client.run_tesseract("jacobian", payload, run_id)
 
     @requires_client
     def jacobian_vector_product(
@@ -375,6 +383,7 @@ def jacobian_vector_product(
         jvp_inputs: list[str],
         jvp_outputs: list[str],
         tangent_vector: dict,
+        run_id: str | None = None,
     ) -> dict:
         """Calculate the Jacobian Vector Product (JVP) of (some of the) outputs w.r.t. (some of the) inputs.
 
@@ -383,6 +392,8 @@ def jacobian_vector_product(
             jvp_inputs: Inputs with respect to which derivatives will be calculated.
             jvp_outputs: Outputs which will be differentiated.
             tangent_vector: Element of the tangent space to multiply with the Jacobian.
+            run_id: a string to identify the run. Run outputs will be located
+                    in a directory suffixed with this id.
 
         Returns:
             dictionary with the results.
@@ -398,7 +409,7 @@ def jacobian_vector_product(
             "jvp_outputs": jvp_outputs,
             "tangent_vector": tangent_vector,
         }
-        return self._client.run_tesseract("jacobian_vector_product", payload)
+        return self._client.run_tesseract("jacobian_vector_product", payload, run_id)
 
     @requires_client
     def vector_jacobian_product(
@@ -407,6 +418,7 @@ def vector_jacobian_product(
         vjp_inputs: list[str],
         vjp_outputs: list[str],
         cotangent_vector: dict,
+        run_id: str | None = None,
     ) -> dict:
         """Calculate the Vector Jacobian Product (VJP) of (some of the) outputs w.r.t. (some of the) inputs.
 
@@ -415,6 +427,8 @@ def vector_jacobian_product(
             vjp_inputs: Inputs with respect to which derivatives will be calculated.
             vjp_outputs: Outputs which will be differentiated.
             cotangent_vector: Element of the cotangent space to multiply with the Jacobian.
+            run_id: a string to identify the run. Run outputs will be located
+                    in a directory suffixed with this id.
 
 
         Returns:
@@ -431,7 +445,7 @@ def vector_jacobian_product(
             "vjp_outputs": vjp_outputs,
             "cotangent_vector": cotangent_vector,
         }
-        return self._client.run_tesseract("vector_jacobian_product", payload)
+        return self._client.run_tesseract("vector_jacobian_product", payload, run_id)
 
 
 def _tree_map(func: Callable, tree: Any, is_leaf: Callable | None = None) -> Any:
@@ -506,7 +520,11 @@ def url(self) -> str:
         return self._url
 
     def _request(
-        self, endpoint: str, method: str = "GET", payload: dict | None = None
+        self,
+        endpoint: str,
+        method: str = "GET",
+        payload: dict | None = None,
+        run_id: str | None = None,
     ) -> dict:
         url = f"{self.url}/{endpoint.lstrip('/')}"
 
@@ -517,7 +535,10 @@ def _request(
         else:
             encoded_payload = None
 
-        response = requests.request(method=method, url=url, json=encoded_payload)
+        params = {"run_id": run_id} if run_id is not None else {}
+        response = requests.request(
+            method=method, url=url, json=encoded_payload, params=params
+        )
 
         if response.status_code == requests.codes.unprocessable_entity:
             # Try and raise a more helpful error if the response is a Pydantic error
@@ -570,12 +591,16 @@ def _request(
 
         return data
 
-    def run_tesseract(self, endpoint: str, payload: dict | None = None) -> dict:
+    def run_tesseract(
+        self, endpoint: str, payload: dict | None = None, run_id: str | None = None
+    ) -> dict:
         """Run a Tesseract endpoint.
 
         Args:
             endpoint: The endpoint to run.
             payload: The payload to send to the endpoint.
+            run_id: a string to identify the run. Run outputs will be located
+                    in a directory suffixed with this id.
 
         Returns:
             The loaded JSON response from the endpoint, with decoded arrays.
@@ -591,7 +616,7 @@ def run_tesseract(self, endpoint: str, payload: dict | None = None) -> dict:
         if endpoint == "openapi_schema":
             endpoint = "openapi.json"
 
-        return self._request(endpoint, method, payload)
+        return self._request(endpoint, method, payload, run_id)
 
 
 class LocalClient:
@@ -606,12 +631,15 @@ def __init__(self, tesseract_api: ModuleType) -> None:
         }
         self._openapi_schema = create_rest_api(tesseract_api).openapi()
 
-    def run_tesseract(self, endpoint: str, payload: dict | None = None) -> dict:
+    def run_tesseract(
+        self, endpoint: str, payload: dict | None = None, run_id: str | None = None
+    ) -> dict:
         """Run a Tesseract endpoint.
 
         Args:
             endpoint: The endpoint to run.
             payload: The payload to send to the endpoint.
+            run_id: a string to identify the run.
 
         Returns:
             The loaded JSON response from the endpoint, with decoded arrays.

tests/sdk_tests/test_tesseract.py

@@ -130,7 +130,11 @@ def test_serve_lifecycle(mock_serving, mock_clients):
                 pass
 
 
-def test_HTTPClient_run_tesseract(mocker):
+@pytest.mark.parametrize(
+    "run_id",
+    [None, "fizzbuzz"],
+)
+def test_HTTPClient_run_tesseract(mocker, run_id):
     mock_response = mocker.Mock()
     mock_response.json.return_value = {"result": [4, 4, 4]}
     mock_response.raise_for_status = mocker.Mock()
@@ -142,13 +146,15 @@ def test_HTTPClient_run_tesseract(mocker):
 
     client = HTTPClient("somehost")
 
-    out = client.run_tesseract("apply", {"inputs": {"a": 1}})
+    out = client.run_tesseract("apply", {"inputs": {"a": 1}}, run_id=run_id)
 
     assert out == {"result": [4, 4, 4]}
+    expected_params = {} if run_id is None else {"run_id": run_id}
     mocked_request.assert_called_with(
         method="POST",
         url="http://somehost/apply",
         json={"inputs": {"a": 1}},
+        params=expected_params,
     )