`start_run` context manager dominates `from_tesseract_api` overhead

[jpbrodrick89]

For the noop Tesseract used for benchmarking >90% (potentially much closer to 100%) of the 0.4ms runtime is dominated by the start_run context manager. This context manager is used to create a file backend for mpa logging and redirecting stdio for setting up the teepiping. The breakdown seems to be in equal parts

• Creating directories
• Writing the csv
• Redirecting stdio

In from_tesseract_api there is very little value in this, typically errors will bubble up through Python and we use results straight away. Some options for dealing with this:

1. Lazy initialization: don't create the FileBackend/logs infrastructure unless someone actually calls log_metric/log_parameter/log_artifact
2. Reuse the backend across calls instead of creating a new one per request
3. Skip stdio redirection when no log_sink is provided and the server is already handling its own logging
4. For served Tesseract setup stdio redirection at serve time not on every request
5. Make MPA opt-in: if the user isn't using metrics/artifacts, skip the whole thing entirely

I'd probably go for a combination of the last three, but would be happy with 1, 3 and 4 instead.

I know its a tiny time but I think it would be really awesome if from_tesseract_api could be more perfomant for mini-Tesseracts (e.g. with runtimes of 100µs or less) in tesseract-jax to take advantage of the convenient Jax primitive wrapping and resulting language interop.

Tagging @linusseelinger and @sakoush for visibility in case you had any opinions if this affects P4D interactions.

[dionhaefner]

I don't necessarily disagree but before we move on this we should make sure that this is an actual bottleneck in envisioned use cases for micro-Tesseracts (premature optimization is... you know what). Concretely, I'd assume internal copies between JAX and NumPy to dominate overheads in Tesseract-JAX, and any optimizations we do in Tesseract Core beyond what we have won't make a measurable impact.

[jpbrodrick89]

Based on quick benchmarks for array sizes up to 100k I can confirm additional overhead is < 12% of the runtime when dtypes match (~0.05ms but hard to cleanly separate). At 100k dtype mismatch casting cost can be noticeable at ~0.1ms (20%). However, at array sizes of 1M (even when dtypes match) the emit_python_callback overhead begins to dominate over the start_run overhead under jit. Therefore, we should still get about 10x speedup by addressing this for medium sized arrays. Note that you also pay the 0.4ms when calling abstract eval (e.g. in eager mode or at trace time).

[jpbrodrick89]

Note for larger arrays jax.experimental.buffer_callback would eliminate the memory callback overhead meaning we really could support mini-Tesseracts operating on large arrays.

[jpbrodrick89]

Note that this overhead was only introduced in version 1.4, before then the current implementation of tesseract-jax was managing 0.04ms performance. This seems to have entirely been caused by #482.

  ┌────────────────────────────┬───────┬───────┬───────┬───────┐
  │         Benchmark          │ 1.3.0 │ 1.4.0 │ 1.5.1 │ 1.6.0 │
  ├────────────────────────────┼───────┼───────┼───────┼───────┤
  │ noop_api/apply_jit [1K]    │ 0.041 │ 0.042 │ 109   │ 0.36  │
  ├────────────────────────────┼───────┼───────┼───────┼───────┤
  │ noop_api/apply_jit [1M]    │ 1.74  │ 1.75  │ 115   │ 0.89  │
  ├────────────────────────────┼───────┼───────┼───────┼───────┤
  │ noop_api/cast_float64 [1K] │ 0.041 │ 0.042 │ 109   │ 0.40  │
  ├────────────────────────────┼───────┼───────┼───────┼───────┤
  │ noop_api/vmap [1K]         │ 22.5  │ 23.1  │ 1,265 │ 28.6  │
  ├────────────────────────────┼───────┼───────┼───────┼───────┤
  │ vectoradd_api/vjp [1K]     │ 1.37  │ 1.52  │ 344   │ 5.58  │
  ├────────────────────────────┼───────┼───────┼───────┼───────┤
  │ vectoradd_api/vjp [10M]    │ 69.3  │ 74.9  │ 391   │ 31.7  │
  └────────────────────────────┴───────┴───────┴───────┴───────┘