Add a sample function that computes a batch of template results in parallel
#417
Conversation
|
Unfortunately, adding type annotations to |
|
I think it would be better to implement this slightly differently, so converting to draft for now. |
effectful/handlers/llm/sampling.py
Outdated
| completed = futures.wait(tasks, return_when=futures.ALL_COMPLETED) | ||
| return [t.result() for t in completed.done] | ||
|
|
||
| return handler({Template.__call__: _template_call})(template) |
There was a problem hiding this comment.
what is the benefit of sample using the handler machinery here rather than directly retrieving the information from the template passed in?
There was a problem hiding this comment.
and then recursively calling Template.__call__ to reuse the existing machinery? also should sample be a defop?
| try: | ||
| result = fwd() | ||
| except Exception as e: | ||
| lock.release() |
There was a problem hiding this comment.
if fwd() doesn't raise an exception this won't release the lock no?
with lock:
return fwd()There was a problem hiding this comment.
also it might be good to have some test for this functionality, mocking the llm, just checking that we're doing concurrency correctly
There was a problem hiding this comment.
if fwd() doesn't raise an exception this won't release the lock no?
with lock: return fwd()
We actually don't want to. If the tool call completes successfully, we keep the lock up to the point that we start another completion. Releasing the lock here allows the remaining threads to proceed as this thread dies. The fwd in the completion handler can fail safely because we don't hold the lock.
This raises a good point though: what should we do when one of the threads fails? Right now we reraise the exception from the failed thread and you get no samples. Maybe instead you should get the successful samples?
| return result | ||
|
|
||
| def _tool_call(*args, **kwargs): | ||
| lock.acquire() |
There was a problem hiding this comment.
I think this call to acquire isn't necessary.
The sampling methods in #415 and #412 could be mostly implemented on top of a parallel
sampleoperation. You get a little less parallelism, since you would wait for all voters to complete before scheduling new ones.