Cache (2/5): Cache inference at do_inference - #122
Conversation
97138e9 to
965bacf
Compare
ErlisLushtaku
changed the title
ErlisLushtaku
changed the title
b1156b8 to
52a4620
Compare
58f0f59 to
969347c
Compare
|
|
||
|
|
||
| @dataclass(frozen=True) | ||
| class InferenceCache(ABC): |
There was a problem hiding this comment.
Should we create a seperate file for storing these cache objects?
There was a problem hiding this comment.
Agreed, the cache-specific types now live in judgearena/cache/inference.py in #127. The provider-neutral InferenceResult stays in judgearena/inference.py.
| cached_rows = store.query(input_hashes).set_index("input_hash") | ||
| results: list[InferenceResult | None] = [ | ||
| ( | ||
| InferenceResult(**cache.cached_result(cached_rows.loc[key]).__dict__) |
There was a problem hiding this comment.
Why we are converting CachedInferenceResult back to InferenceResult? If so we may not need CachedInferenceResultat all
There was a problem hiding this comment.
Yes, CachedInferenceResult is removed and cache hits now return InferenceResult directly.
| ) | ||
| for index, result in zip(missing_indices, generated, strict=True): | ||
| results[index] = result | ||
| cache.save_outputs( |
There was a problem hiding this comment.
Can this fail? If it fails for some reason the entire pipeline fails. We can put try except
try:
cache.save_outputs(...)
except CacheWriteError as exc:
logger.warning("Could not save inference cache: %s", exc)
There was a problem hiding this comment.
Good catch, cache write failures now log and preserve the generated results.
| return_top_logprobs: bool = False, | ||
| *, | ||
| stage: str = "unspecified", | ||
| cache_metadata: list[dict] | None = None, |
There was a problem hiding this comment.
I think one thing that is confusing for me is this PR introduces some changes like cache_metadata which will be provided in the future PR's in #124. Design-wise it makes sense to first introduce it without wiring however while reading the PR it creates some readability problems. I think if necessary, creating large PR's (~1000-2000LOC) where we introduce the functionality while introducing the changes makes more sense to me.
Not a requirement for this PR (as I will review the entire stack as a single one)
|
|
||
|
|
||
| def do_inference( | ||
| def _do_inference_uncached( |
There was a problem hiding this comment.
We can rename this to _call_model or _invoke_model or similar (I am not creative enough). Because we will be using it to do inference for uncached samples, this reads like we will be doing inference without using the cache (although it is true we are doing it for uncached inputs, flow can be better)
There was a problem hiding this comment.
Renamed it to _run_backend_inference.
| stage=stage, | ||
| ) | ||
|
|
||
| try: |
There was a problem hiding this comment.
I know we need this for closing the cache but it becomes hard to parse with lots of try: excepts
There was a problem hiding this comment.
Agreed, I moved cache open and close handling into a context manager and read and write failures into small helpers.
2 participants
Description
Moves cache lookup and writes to
do_inference.PreparedModeldefers backend construction until an input misses the cache.This is stacked on #121.