teel 0.2.2__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- teel-0.2.2/.github/CODEOWNERS +1 -0
- teel-0.2.2/.github/settings.yml +52 -0
- teel-0.2.2/.github/workflows/ci.yml +28 -0
- teel-0.2.2/.github/workflows/release.yml +35 -0
- teel-0.2.2/.gitignore +39 -0
- teel-0.2.2/PKG-INFO +205 -0
- teel-0.2.2/README.md +189 -0
- teel-0.2.2/examples/code_review.py +150 -0
- teel-0.2.2/examples/hello_world.py +55 -0
- teel-0.2.2/examples/tree_of_thoughts.py +64 -0
- teel-0.2.2/pyproject.toml +51 -0
- teel-0.2.2/setup.cfg +4 -0
- teel-0.2.2/setup.py +11 -0
- teel-0.2.2/teel/__init__.py +168 -0
- teel-0.2.2/teel/frame.py +228 -0
- teel-0.2.2/teel/frameinspect.c +145 -0
- teel-0.2.2/teel/lookahead.py +781 -0
- teel-0.2.2/teel/py.typed +0 -0
- teel-0.2.2/teel/shadow_copy.py +83 -0
- teel-0.2.2/teel/timeline.py +121 -0
- teel-0.2.2/teel.egg-info/PKG-INFO +205 -0
- teel-0.2.2/teel.egg-info/SOURCES.txt +25 -0
- teel-0.2.2/teel.egg-info/dependency_links.txt +1 -0
- teel-0.2.2/teel.egg-info/scm_file_list.json +24 -0
- teel-0.2.2/teel.egg-info/top_level.txt +1 -0
- teel-0.2.2/tests/__init__.py +0 -0
- teel-0.2.2/tests/test_lookahead.py +500 -0
|
@@ -0,0 +1 @@
|
|
|
1
|
+
* @shahar-lev @sternby
|
|
@@ -0,0 +1,52 @@
|
|
|
1
|
+
# These settings are synced to GitHub by https://probot.github.io/apps/settings/
|
|
2
|
+
|
|
3
|
+
repository:
|
|
4
|
+
description: "🚀 Teel: Speculative ahead-of-time parallelization of plain synchronous Python via a shadow bytecode interpreter"
|
|
5
|
+
topics: python, concurrency, parallelism, speculative-execution, bytecode, cpython
|
|
6
|
+
|
|
7
|
+
has_issues: true
|
|
8
|
+
has_projects: false
|
|
9
|
+
has_wiki: false
|
|
10
|
+
has_downloads: false
|
|
11
|
+
|
|
12
|
+
default_branch: main
|
|
13
|
+
|
|
14
|
+
allow_squash_merge: true
|
|
15
|
+
allow_merge_commit: true
|
|
16
|
+
allow_rebase_merge: true
|
|
17
|
+
allow_auto_merge: false
|
|
18
|
+
delete_branch_on_merge: true
|
|
19
|
+
|
|
20
|
+
enable_automated_security_fixes: true
|
|
21
|
+
enable_vulnerability_alerts: true
|
|
22
|
+
|
|
23
|
+
archived: false
|
|
24
|
+
|
|
25
|
+
labels:
|
|
26
|
+
- name: bug
|
|
27
|
+
color: '#cc0000'
|
|
28
|
+
description: An issue with the system.
|
|
29
|
+
- name: feature
|
|
30
|
+
color: '#336699'
|
|
31
|
+
description: New functionality
|
|
32
|
+
- name: WIP
|
|
33
|
+
color: '#faf064'
|
|
34
|
+
description: Work in progress
|
|
35
|
+
- name: draft
|
|
36
|
+
color: '#c800c8'
|
|
37
|
+
|
|
38
|
+
branches:
|
|
39
|
+
- name: main
|
|
40
|
+
protection:
|
|
41
|
+
required_pull_request_reviews:
|
|
42
|
+
required_approving_review_count: 1
|
|
43
|
+
dismiss_stale_reviews: true
|
|
44
|
+
require_code_owner_reviews: true
|
|
45
|
+
dismissal_restrictions: {}
|
|
46
|
+
required_status_checks:
|
|
47
|
+
strict: true
|
|
48
|
+
contexts:
|
|
49
|
+
- test
|
|
50
|
+
enforce_admins: true
|
|
51
|
+
restrictions: null
|
|
52
|
+
required_conversation_resolution: true
|
|
@@ -0,0 +1,28 @@
|
|
|
1
|
+
name: CI
|
|
2
|
+
|
|
3
|
+
on:
|
|
4
|
+
push:
|
|
5
|
+
branches: [main]
|
|
6
|
+
pull_request:
|
|
7
|
+
branches: [main]
|
|
8
|
+
|
|
9
|
+
concurrency:
|
|
10
|
+
group: ci-${{ github.ref }}
|
|
11
|
+
cancel-in-progress: true
|
|
12
|
+
|
|
13
|
+
jobs:
|
|
14
|
+
test:
|
|
15
|
+
runs-on: ubuntu-latest
|
|
16
|
+
steps:
|
|
17
|
+
- uses: actions/checkout@v4
|
|
18
|
+
|
|
19
|
+
- name: Install uv
|
|
20
|
+
uses: astral-sh/setup-uv@v6
|
|
21
|
+
with:
|
|
22
|
+
enable-cache: true
|
|
23
|
+
|
|
24
|
+
- name: Sync environment
|
|
25
|
+
run: uv sync
|
|
26
|
+
|
|
27
|
+
- name: Run tests
|
|
28
|
+
run: uv run python -m unittest discover -s tests -v
|
|
@@ -0,0 +1,35 @@
|
|
|
1
|
+
name: Release
|
|
2
|
+
|
|
3
|
+
on:
|
|
4
|
+
push:
|
|
5
|
+
branches: [main]
|
|
6
|
+
|
|
7
|
+
permissions:
|
|
8
|
+
contents: write
|
|
9
|
+
id-token: write
|
|
10
|
+
|
|
11
|
+
concurrency:
|
|
12
|
+
group: release
|
|
13
|
+
cancel-in-progress: false
|
|
14
|
+
|
|
15
|
+
jobs:
|
|
16
|
+
release:
|
|
17
|
+
runs-on: ubuntu-latest
|
|
18
|
+
environment: pypi
|
|
19
|
+
steps:
|
|
20
|
+
- uses: actions/checkout@v4
|
|
21
|
+
with:
|
|
22
|
+
fetch-depth: 0
|
|
23
|
+
|
|
24
|
+
- name: Install uv
|
|
25
|
+
uses: astral-sh/setup-uv@v6
|
|
26
|
+
|
|
27
|
+
- name: Semantic release
|
|
28
|
+
id: release
|
|
29
|
+
env:
|
|
30
|
+
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
|
31
|
+
run: uvx --from "python-semantic-release>=10.5.3" semantic-release version --no-commit
|
|
32
|
+
|
|
33
|
+
- name: Publish to PyPI
|
|
34
|
+
if: steps.release.outputs.released == 'true'
|
|
35
|
+
run: uv publish --trusted-publishing always
|
teel-0.2.2/.gitignore
ADDED
|
@@ -0,0 +1,39 @@
|
|
|
1
|
+
# Byte-compiled / optimized / DLL files
|
|
2
|
+
__pycache__/
|
|
3
|
+
*.py[cod]
|
|
4
|
+
*$py.class
|
|
5
|
+
|
|
6
|
+
# C extensions
|
|
7
|
+
*.so
|
|
8
|
+
|
|
9
|
+
# Distribution / packaging
|
|
10
|
+
build/
|
|
11
|
+
dist/
|
|
12
|
+
*.egg-info/
|
|
13
|
+
*.egg
|
|
14
|
+
|
|
15
|
+
# Unit test / coverage reports
|
|
16
|
+
htmlcov/
|
|
17
|
+
.coverage
|
|
18
|
+
.coverage.*
|
|
19
|
+
.cache
|
|
20
|
+
coverage.xml
|
|
21
|
+
*.cover
|
|
22
|
+
.hypothesis/
|
|
23
|
+
.pytest_cache/
|
|
24
|
+
|
|
25
|
+
# Environments
|
|
26
|
+
.env
|
|
27
|
+
.venv
|
|
28
|
+
env/
|
|
29
|
+
venv/
|
|
30
|
+
ENV/
|
|
31
|
+
|
|
32
|
+
# IDEs
|
|
33
|
+
.idea/
|
|
34
|
+
.vscode/
|
|
35
|
+
*.swp
|
|
36
|
+
|
|
37
|
+
# Claude Code local state (per-developer)
|
|
38
|
+
/.claude/worktrees/
|
|
39
|
+
/.claude/settings.local.json
|
teel-0.2.2/PKG-INFO
ADDED
|
@@ -0,0 +1,205 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: teel
|
|
3
|
+
Version: 0.2.2
|
|
4
|
+
Summary: Speculative ahead-of-time parallelization of plain synchronous Python via a shadow bytecode interpreter
|
|
5
|
+
Author: AI21
|
|
6
|
+
Project-URL: Repository, https://github.com/AI21Labs/teel
|
|
7
|
+
Project-URL: Issues, https://github.com/AI21Labs/teel/issues
|
|
8
|
+
Keywords: concurrency,parallelism,speculative-execution,lookahead,bytecode,cache,threadpool,cpython
|
|
9
|
+
Classifier: Development Status :: 3 - Alpha
|
|
10
|
+
Classifier: Intended Audience :: Developers
|
|
11
|
+
Classifier: Programming Language :: Python :: 3
|
|
12
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
13
|
+
Classifier: Topic :: Software Development :: Libraries
|
|
14
|
+
Requires-Python: ==3.12.13
|
|
15
|
+
Description-Content-Type: text/markdown
|
|
16
|
+
|
|
17
|
+
# 🚀 Teel - Turbo-Efficient Execution Library
|
|
18
|
+
|
|
19
|
+
### Async magic for synchronous code
|
|
20
|
+
|
|
21
|
+
Since LLMs and Agents, a big portion of production code became IO-bound.
|
|
22
|
+
|
|
23
|
+
A lot of approaches to writing faster code involve introducing concurrent concepts like futures or coroutines, trading off simplicity for performance.
|
|
24
|
+
Before you know it, your code is filled with concurrency bits like `.result()` or `asyncio.gather`.
|
|
25
|
+
|
|
26
|
+
We wanted to stop the async "mechanicals" of ruining the simplicity of our main logic.
|
|
27
|
+
|
|
28
|
+
## How it works
|
|
29
|
+
|
|
30
|
+
Consider the following simple snippet:
|
|
31
|
+
|
|
32
|
+
```python
|
|
33
|
+
y = f1(x)
|
|
34
|
+
z = f2(x)
|
|
35
|
+
w = f3(y, z)
|
|
36
|
+
```
|
|
37
|
+
|
|
38
|
+
Assume each function `f1, f2, f3` takes about a second to complete.
|
|
39
|
+
In a normal run, when everything runs sequentially, this code takes about 3 seconds to complete.
|
|
40
|
+
|
|
41
|
+
In an optimized run, the runtime magically realizes that the `f2(x)` is on the horizon and runs it in parallel.
|
|
42
|
+
When `f1(x)` finishes, `f2(x)` is magically ready as well - so we immediately execute `f3(y, z)`.
|
|
43
|
+
|
|
44
|
+
**A 33% of the time was saved 🚀**
|
|
45
|
+
|
|
46
|
+
That's basically the deal -
|
|
47
|
+
Teel adds the necessary magic to make sequential code run optimized by submitting these functions to be run on multiple threads *ahead of time*.
|
|
48
|
+
This lookahead mechanism does not have to be perfect - it operates on a best effort basis, sniffing out possible future calls to cache-able functions.
|
|
49
|
+
The hard part is to make sure this does not mutate state tied to the MainThread run, otherwise we risk corrupting the execution.
|
|
50
|
+
|
|
51
|
+
### Speed-up Results
|
|
52
|
+
|
|
53
|
+
| Code | Original E2E duration | 🚀 Teel E2E duration | Speed-up |
|
|
54
|
+
|-------|----------------------|--------------------------|----------|
|
|
55
|
+
| [tree-of-thoughts-llm](https://github.com/princeton-nlp/tree-of-thought-llm) | 307s | 38s | 8.07x |
|
|
56
|
+
| RLM inside LangChain's Deep Agents (python) | 150s | 58s | 2.58x |
|
|
57
|
+
| RLM inside LangChain's Deep Agents (async JS) | 71s | 58s | 1.22x |
|
|
58
|
+
|
|
59
|
+
|
|
60
|
+
## Getting started with Teel
|
|
61
|
+
|
|
62
|
+
### 1. Install
|
|
63
|
+
Clone the repo
|
|
64
|
+
```
|
|
65
|
+
git clone https://github.com/AI21Labs/teel.git
|
|
66
|
+
```
|
|
67
|
+
|
|
68
|
+
Install dependencies and run hello world:
|
|
69
|
+
```
|
|
70
|
+
uv sync
|
|
71
|
+
uv run python -m examples.hello_world
|
|
72
|
+
```
|
|
73
|
+
|
|
74
|
+
The `uv sync` command takes care of building the `frameinspect` C extension.
|
|
75
|
+
The `hello_world` example runs an insanely simple diamond pipeline with and without `Teel`, prints time metrics and saves a timeline into `teel-trace.json`, available for viewing at [ui.perfetto.dev](https://ui.perfetto.dev/).
|
|
76
|
+
|
|
77
|
+
|
|
78
|
+
>Note that we currently have a hard constraint on exact Python version 3.12.13.
|
|
79
|
+
>This is because `frameinspect.c` needs CPython's internal headers, so if you're running with a different Python version you may need to adjust the extension.
|
|
80
|
+
|
|
81
|
+
|
|
82
|
+
### 2. Turbo-charge your code
|
|
83
|
+
|
|
84
|
+
The API is thin - just three decorators and one context manager.
|
|
85
|
+
|
|
86
|
+
1. Mark your "heavy" functions (the ones that are cache-able and you wish to auto-parallelize) with `@cacheable`.
|
|
87
|
+
The cache key for these calls will be based on a `json.dumps` of the arguments (so don't throw in there any funky stuff).
|
|
88
|
+
```python
|
|
89
|
+
@cacheable
|
|
90
|
+
def heavy_calc(x, y, z):
|
|
91
|
+
...
|
|
92
|
+
```
|
|
93
|
+
|
|
94
|
+
2. Mark the main functions you want Teel to look for potential optimizations with `@jit`.
|
|
95
|
+
```python
|
|
96
|
+
@jit
|
|
97
|
+
def main():
|
|
98
|
+
result = heavy_calc(a, b, c)
|
|
99
|
+
...
|
|
100
|
+
```
|
|
101
|
+
|
|
102
|
+
3. Optionally, decorate fast helper functions (like string formatting or dict access) with `@inline`. this helps adding more stuff to be submitted ahead-of-time.
|
|
103
|
+
```python
|
|
104
|
+
@inline
|
|
105
|
+
def get_user_data(raw_user):
|
|
106
|
+
...
|
|
107
|
+
return raw_user.get('data', {}).get('result', None)
|
|
108
|
+
```
|
|
109
|
+
|
|
110
|
+
4. Wrap everything with a `with Teel(): ...` to activate the magic (otherwise it just runs synchronously).
|
|
111
|
+
```python
|
|
112
|
+
with Teel():
|
|
113
|
+
main()
|
|
114
|
+
```
|
|
115
|
+
|
|
116
|
+
>Because Teel dispatches `@cacheable` functions ahead-of-time, these functions must be **pure and idempotent**, meaning they give the same output for same input, with no side effects. That means nothing that writes to a DB or randomly selects a return value.
|
|
117
|
+
|
|
118
|
+
### 3. Tracing
|
|
119
|
+
|
|
120
|
+
You can provide `Teel()` with a `trace_path=` and it will record a timeline of `@cacheable`-decorated execution, so you'll be able to see just what exactly ran in parallel.
|
|
121
|
+
The output file is saved in a Chrome trace-event json which can be opened at [ui.perfetto.dev](https://ui.perfetto.dev/) to view the timeline.
|
|
122
|
+
In addition to start/end times and durations, you can see in the UI that each call has its own index (`callno`) and the index of the call it originated from (`parent_callno`).
|
|
123
|
+
In the example above, the `f2` call is the parent of the `f3` call (because the `f3` call was discovered in a lookahead search conducted at the `f2` point).
|
|
124
|
+
|
|
125
|
+
## Under the hood
|
|
126
|
+
|
|
127
|
+
Teel's **lookahead mechanism** starts by duplicating the execution state into a limited "shadow interpreter".
|
|
128
|
+
There it uses `sys._getframe()` to retrieve the frame stack, and applies a custom-written C extension that retrieves the frame's full value stack.
|
|
129
|
+
It uses `f_lasti` of each frame to know the last bytecode instruction that was executed.
|
|
130
|
+
|
|
131
|
+
The shadow interpreter iterates over bytecode opcodes, processing them and mutating its own shadow value stack per frame. Think of it as a (obviously slower) reimplementation of CPython's loop that can handle pending (not yet evaluated) values. It cautiously does copy-on-write as to not screw with the state of the main thread.
|
|
132
|
+
|
|
133
|
+
The shadow interpreter may:
|
|
134
|
+
|
|
135
|
+
1. Step into (user decorated) functions and shadow-execute them as well.
|
|
136
|
+
2. Skip instructions/calls that cannot yet be executed - marking them as "pending".
|
|
137
|
+
3. Run (by regular invocation) certain (user decorated) inline functions that it deems safe to run (not going to corrupt the main state).
|
|
138
|
+
4. Perform copy-on-write on mutating instructions, further maintaining a shadow state while keeping the main state intact.
|
|
139
|
+
5. Perform speculative execution by cloning state and following conditional jump instructions (up to configured width/depth to avoid exponential explosion).
|
|
140
|
+
|
|
141
|
+
The shadow interpreter submits function calls to be executed by ThreadPoolExecutor in the background. This makes the cache-able functions warmed up before MainThread reaches them, thus reducing runtime.
|
|
142
|
+
|
|
143
|
+
These background calls, when they finish, can also trigger more searches and thus more background calls (until MainThread reaches another point in the execution where the previous generation of pending background calls is discarded in favor of a new one).
|
|
144
|
+
This ability of piercing into (multiple potential) futures enables Teel to infer runtime parallelism by deducing dependencies and pre-invoking functions on speculative execution and thus reducing runtime.
|
|
145
|
+
|
|
146
|
+
### Caching calls by their inputs
|
|
147
|
+
|
|
148
|
+
Lookahead search can be performed from multiple threads and places and so the same (with equal arguments) function call can be discovered (and submitted) multiple times.
|
|
149
|
+
Teel keeps a dict of pending/done tasks, keyed by the function itself and `json.dumps` of the input arguments, in order to execute just one instance of such call+arguments pair.
|
|
150
|
+
Worth noting that if the owner's call failed, Teel does not cache the failure and assumes it's retryable (might be worth revisiting this notion and retry only for specific exceptions).
|
|
151
|
+
|
|
152
|
+
## Tests
|
|
153
|
+
|
|
154
|
+
The small `tests/test_lookahead.py` (desperately needs to grow in size) helps with both making sure the search can agree with the CPython interpreter about the return value of a function and more importantly that it doesn't mutate any state that might affect the CPython outcome.
|
|
155
|
+
|
|
156
|
+
## Examples
|
|
157
|
+
|
|
158
|
+
### Tree of Thoughts
|
|
159
|
+
|
|
160
|
+
We took [Princeton's `tree-of-thoughts-llm`](https://github.com/princeton-nlp/tree-of-thought-llm), unmodified, and ran several samples of the "Game of 24" dataset with GPT-4.
|
|
161
|
+
|
|
162
|
+
|
|
163
|
+
We applied our decorators outside of the original code, decorating `@jit`, `@cacheable` and `@inline`.
|
|
164
|
+
|
|
165
|
+
|
|
166
|
+
Because the original is written in naive synchronous code we managed to get an average reduction of 83% in runtime compared to regular execution (from as low as 148s optimized into 24s and as high as 307s optimized into 38s).
|
|
167
|
+
|
|
168
|
+
See [`examples/tree_of_thoughts.py`](examples/tree_of_thoughts.py).
|
|
169
|
+
|
|
170
|
+
### Recursive Language Models in Deep Agents
|
|
171
|
+
|
|
172
|
+
We took LangChain's Deep Agents (specifically the Deep Agents Code CLI, or `dcode`) for a spin.
|
|
173
|
+
|
|
174
|
+
|
|
175
|
+
See [How to Use RLMs in Deep Agents](https://www.langchain.com/blog/how-to-use-rlms-in-deep-agents).
|
|
176
|
+
Our [Deep Agents fork](https://github.com/AI21Labs/deepagents) uses Teel.
|
|
177
|
+
|
|
178
|
+
The `dcode` CLI comes with an RLM-based `js_eval` tool - the model can choose to use it to run a JavaScript snippet. The snippet *itself* can contain tool calls (such as `tools.web_search()`) and also a call to `task()` which runs a subagent (with its own context) as part of the snippet execution, achieving the recursion.
|
|
179
|
+
|
|
180
|
+
The snippet is run in a sandboxed environment using QuickJS.
|
|
181
|
+
|
|
182
|
+
Putting isolation aside, we wanted to measure the effect on the runtime if we replaced the (async) JS with synchronous Python, freeing the LLM from async considerations in writing the code.
|
|
183
|
+
|
|
184
|
+
We ran several prompts whose solutions involve a mixture of web search and LLM calls to be solved end-to-end by a single snippet.
|
|
185
|
+
Compared to synchronous Python we can see substantial improvement.
|
|
186
|
+
As for the JS comparison - we noticed a slight improvement.
|
|
187
|
+
|
|
188
|
+
Taken with a grain of salt (because even the same prompt produces different JS and Python solutions), we can point out that not all JS snippets use `Promise.all` where they should (instead of just `await`) and that JS doesn't bother pre-running fallback logic ("do B if A was unsatisfactory"), unlike `Teel` which speculatively runs these fallbacks ahead of time.
|
|
189
|
+
|
|
190
|
+
Running 10 prompts as described, on average Teel takes 58s, original async JS takes 71s and plain synchronous Python takes 150s.
|
|
191
|
+
|
|
192
|
+
This is 61% faster than plain synchronous Python and 19% faster than the JS baseline.
|
|
193
|
+
|
|
194
|
+
## GIL peskiness
|
|
195
|
+
|
|
196
|
+
In the current implementation we ship background calls to a `ThreadPoolExecutor`.
|
|
197
|
+
If your `@cacheable` does I/O or network work (like `requests.get`) or GIL-free CPU work (`numpy`) then you're good to go.
|
|
198
|
+
If your `@cacheable` is CPU-bound and also GIL-constrained then no dice.
|
|
199
|
+
It's not really a difficult hurdle to implement - either by switching to `ProcessPoolExecutor` or
|
|
200
|
+
`InterpreterPoolExecutor` (but that's been added only in Python 3.14), but it's not currently an option in the code.
|
|
201
|
+
|
|
202
|
+
## Disclaimer
|
|
203
|
+
|
|
204
|
+
This is a Proof-of-Concept.
|
|
205
|
+
Further work needs to be done to adjust this for other CPython versions.
|
teel-0.2.2/README.md
ADDED
|
@@ -0,0 +1,189 @@
|
|
|
1
|
+
# 🚀 Teel - Turbo-Efficient Execution Library
|
|
2
|
+
|
|
3
|
+
### Async magic for synchronous code
|
|
4
|
+
|
|
5
|
+
Since LLMs and Agents, a big portion of production code became IO-bound.
|
|
6
|
+
|
|
7
|
+
A lot of approaches to writing faster code involve introducing concurrent concepts like futures or coroutines, trading off simplicity for performance.
|
|
8
|
+
Before you know it, your code is filled with concurrency bits like `.result()` or `asyncio.gather`.
|
|
9
|
+
|
|
10
|
+
We wanted to stop the async "mechanicals" of ruining the simplicity of our main logic.
|
|
11
|
+
|
|
12
|
+
## How it works
|
|
13
|
+
|
|
14
|
+
Consider the following simple snippet:
|
|
15
|
+
|
|
16
|
+
```python
|
|
17
|
+
y = f1(x)
|
|
18
|
+
z = f2(x)
|
|
19
|
+
w = f3(y, z)
|
|
20
|
+
```
|
|
21
|
+
|
|
22
|
+
Assume each function `f1, f2, f3` takes about a second to complete.
|
|
23
|
+
In a normal run, when everything runs sequentially, this code takes about 3 seconds to complete.
|
|
24
|
+
|
|
25
|
+
In an optimized run, the runtime magically realizes that the `f2(x)` is on the horizon and runs it in parallel.
|
|
26
|
+
When `f1(x)` finishes, `f2(x)` is magically ready as well - so we immediately execute `f3(y, z)`.
|
|
27
|
+
|
|
28
|
+
**A 33% of the time was saved 🚀**
|
|
29
|
+
|
|
30
|
+
That's basically the deal -
|
|
31
|
+
Teel adds the necessary magic to make sequential code run optimized by submitting these functions to be run on multiple threads *ahead of time*.
|
|
32
|
+
This lookahead mechanism does not have to be perfect - it operates on a best effort basis, sniffing out possible future calls to cache-able functions.
|
|
33
|
+
The hard part is to make sure this does not mutate state tied to the MainThread run, otherwise we risk corrupting the execution.
|
|
34
|
+
|
|
35
|
+
### Speed-up Results
|
|
36
|
+
|
|
37
|
+
| Code | Original E2E duration | 🚀 Teel E2E duration | Speed-up |
|
|
38
|
+
|-------|----------------------|--------------------------|----------|
|
|
39
|
+
| [tree-of-thoughts-llm](https://github.com/princeton-nlp/tree-of-thought-llm) | 307s | 38s | 8.07x |
|
|
40
|
+
| RLM inside LangChain's Deep Agents (python) | 150s | 58s | 2.58x |
|
|
41
|
+
| RLM inside LangChain's Deep Agents (async JS) | 71s | 58s | 1.22x |
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
## Getting started with Teel
|
|
45
|
+
|
|
46
|
+
### 1. Install
|
|
47
|
+
Clone the repo
|
|
48
|
+
```
|
|
49
|
+
git clone https://github.com/AI21Labs/teel.git
|
|
50
|
+
```
|
|
51
|
+
|
|
52
|
+
Install dependencies and run hello world:
|
|
53
|
+
```
|
|
54
|
+
uv sync
|
|
55
|
+
uv run python -m examples.hello_world
|
|
56
|
+
```
|
|
57
|
+
|
|
58
|
+
The `uv sync` command takes care of building the `frameinspect` C extension.
|
|
59
|
+
The `hello_world` example runs an insanely simple diamond pipeline with and without `Teel`, prints time metrics and saves a timeline into `teel-trace.json`, available for viewing at [ui.perfetto.dev](https://ui.perfetto.dev/).
|
|
60
|
+
|
|
61
|
+
|
|
62
|
+
>Note that we currently have a hard constraint on exact Python version 3.12.13.
|
|
63
|
+
>This is because `frameinspect.c` needs CPython's internal headers, so if you're running with a different Python version you may need to adjust the extension.
|
|
64
|
+
|
|
65
|
+
|
|
66
|
+
### 2. Turbo-charge your code
|
|
67
|
+
|
|
68
|
+
The API is thin - just three decorators and one context manager.
|
|
69
|
+
|
|
70
|
+
1. Mark your "heavy" functions (the ones that are cache-able and you wish to auto-parallelize) with `@cacheable`.
|
|
71
|
+
The cache key for these calls will be based on a `json.dumps` of the arguments (so don't throw in there any funky stuff).
|
|
72
|
+
```python
|
|
73
|
+
@cacheable
|
|
74
|
+
def heavy_calc(x, y, z):
|
|
75
|
+
...
|
|
76
|
+
```
|
|
77
|
+
|
|
78
|
+
2. Mark the main functions you want Teel to look for potential optimizations with `@jit`.
|
|
79
|
+
```python
|
|
80
|
+
@jit
|
|
81
|
+
def main():
|
|
82
|
+
result = heavy_calc(a, b, c)
|
|
83
|
+
...
|
|
84
|
+
```
|
|
85
|
+
|
|
86
|
+
3. Optionally, decorate fast helper functions (like string formatting or dict access) with `@inline`. this helps adding more stuff to be submitted ahead-of-time.
|
|
87
|
+
```python
|
|
88
|
+
@inline
|
|
89
|
+
def get_user_data(raw_user):
|
|
90
|
+
...
|
|
91
|
+
return raw_user.get('data', {}).get('result', None)
|
|
92
|
+
```
|
|
93
|
+
|
|
94
|
+
4. Wrap everything with a `with Teel(): ...` to activate the magic (otherwise it just runs synchronously).
|
|
95
|
+
```python
|
|
96
|
+
with Teel():
|
|
97
|
+
main()
|
|
98
|
+
```
|
|
99
|
+
|
|
100
|
+
>Because Teel dispatches `@cacheable` functions ahead-of-time, these functions must be **pure and idempotent**, meaning they give the same output for same input, with no side effects. That means nothing that writes to a DB or randomly selects a return value.
|
|
101
|
+
|
|
102
|
+
### 3. Tracing
|
|
103
|
+
|
|
104
|
+
You can provide `Teel()` with a `trace_path=` and it will record a timeline of `@cacheable`-decorated execution, so you'll be able to see just what exactly ran in parallel.
|
|
105
|
+
The output file is saved in a Chrome trace-event json which can be opened at [ui.perfetto.dev](https://ui.perfetto.dev/) to view the timeline.
|
|
106
|
+
In addition to start/end times and durations, you can see in the UI that each call has its own index (`callno`) and the index of the call it originated from (`parent_callno`).
|
|
107
|
+
In the example above, the `f2` call is the parent of the `f3` call (because the `f3` call was discovered in a lookahead search conducted at the `f2` point).
|
|
108
|
+
|
|
109
|
+
## Under the hood
|
|
110
|
+
|
|
111
|
+
Teel's **lookahead mechanism** starts by duplicating the execution state into a limited "shadow interpreter".
|
|
112
|
+
There it uses `sys._getframe()` to retrieve the frame stack, and applies a custom-written C extension that retrieves the frame's full value stack.
|
|
113
|
+
It uses `f_lasti` of each frame to know the last bytecode instruction that was executed.
|
|
114
|
+
|
|
115
|
+
The shadow interpreter iterates over bytecode opcodes, processing them and mutating its own shadow value stack per frame. Think of it as a (obviously slower) reimplementation of CPython's loop that can handle pending (not yet evaluated) values. It cautiously does copy-on-write as to not screw with the state of the main thread.
|
|
116
|
+
|
|
117
|
+
The shadow interpreter may:
|
|
118
|
+
|
|
119
|
+
1. Step into (user decorated) functions and shadow-execute them as well.
|
|
120
|
+
2. Skip instructions/calls that cannot yet be executed - marking them as "pending".
|
|
121
|
+
3. Run (by regular invocation) certain (user decorated) inline functions that it deems safe to run (not going to corrupt the main state).
|
|
122
|
+
4. Perform copy-on-write on mutating instructions, further maintaining a shadow state while keeping the main state intact.
|
|
123
|
+
5. Perform speculative execution by cloning state and following conditional jump instructions (up to configured width/depth to avoid exponential explosion).
|
|
124
|
+
|
|
125
|
+
The shadow interpreter submits function calls to be executed by ThreadPoolExecutor in the background. This makes the cache-able functions warmed up before MainThread reaches them, thus reducing runtime.
|
|
126
|
+
|
|
127
|
+
These background calls, when they finish, can also trigger more searches and thus more background calls (until MainThread reaches another point in the execution where the previous generation of pending background calls is discarded in favor of a new one).
|
|
128
|
+
This ability of piercing into (multiple potential) futures enables Teel to infer runtime parallelism by deducing dependencies and pre-invoking functions on speculative execution and thus reducing runtime.
|
|
129
|
+
|
|
130
|
+
### Caching calls by their inputs
|
|
131
|
+
|
|
132
|
+
Lookahead search can be performed from multiple threads and places and so the same (with equal arguments) function call can be discovered (and submitted) multiple times.
|
|
133
|
+
Teel keeps a dict of pending/done tasks, keyed by the function itself and `json.dumps` of the input arguments, in order to execute just one instance of such call+arguments pair.
|
|
134
|
+
Worth noting that if the owner's call failed, Teel does not cache the failure and assumes it's retryable (might be worth revisiting this notion and retry only for specific exceptions).
|
|
135
|
+
|
|
136
|
+
## Tests
|
|
137
|
+
|
|
138
|
+
The small `tests/test_lookahead.py` (desperately needs to grow in size) helps with both making sure the search can agree with the CPython interpreter about the return value of a function and more importantly that it doesn't mutate any state that might affect the CPython outcome.
|
|
139
|
+
|
|
140
|
+
## Examples
|
|
141
|
+
|
|
142
|
+
### Tree of Thoughts
|
|
143
|
+
|
|
144
|
+
We took [Princeton's `tree-of-thoughts-llm`](https://github.com/princeton-nlp/tree-of-thought-llm), unmodified, and ran several samples of the "Game of 24" dataset with GPT-4.
|
|
145
|
+
|
|
146
|
+
|
|
147
|
+
We applied our decorators outside of the original code, decorating `@jit`, `@cacheable` and `@inline`.
|
|
148
|
+
|
|
149
|
+
|
|
150
|
+
Because the original is written in naive synchronous code we managed to get an average reduction of 83% in runtime compared to regular execution (from as low as 148s optimized into 24s and as high as 307s optimized into 38s).
|
|
151
|
+
|
|
152
|
+
See [`examples/tree_of_thoughts.py`](examples/tree_of_thoughts.py).
|
|
153
|
+
|
|
154
|
+
### Recursive Language Models in Deep Agents
|
|
155
|
+
|
|
156
|
+
We took LangChain's Deep Agents (specifically the Deep Agents Code CLI, or `dcode`) for a spin.
|
|
157
|
+
|
|
158
|
+
|
|
159
|
+
See [How to Use RLMs in Deep Agents](https://www.langchain.com/blog/how-to-use-rlms-in-deep-agents).
|
|
160
|
+
Our [Deep Agents fork](https://github.com/AI21Labs/deepagents) uses Teel.
|
|
161
|
+
|
|
162
|
+
The `dcode` CLI comes with an RLM-based `js_eval` tool - the model can choose to use it to run a JavaScript snippet. The snippet *itself* can contain tool calls (such as `tools.web_search()`) and also a call to `task()` which runs a subagent (with its own context) as part of the snippet execution, achieving the recursion.
|
|
163
|
+
|
|
164
|
+
The snippet is run in a sandboxed environment using QuickJS.
|
|
165
|
+
|
|
166
|
+
Putting isolation aside, we wanted to measure the effect on the runtime if we replaced the (async) JS with synchronous Python, freeing the LLM from async considerations in writing the code.
|
|
167
|
+
|
|
168
|
+
We ran several prompts whose solutions involve a mixture of web search and LLM calls to be solved end-to-end by a single snippet.
|
|
169
|
+
Compared to synchronous Python we can see substantial improvement.
|
|
170
|
+
As for the JS comparison - we noticed a slight improvement.
|
|
171
|
+
|
|
172
|
+
Taken with a grain of salt (because even the same prompt produces different JS and Python solutions), we can point out that not all JS snippets use `Promise.all` where they should (instead of just `await`) and that JS doesn't bother pre-running fallback logic ("do B if A was unsatisfactory"), unlike `Teel` which speculatively runs these fallbacks ahead of time.
|
|
173
|
+
|
|
174
|
+
Running 10 prompts as described, on average Teel takes 58s, original async JS takes 71s and plain synchronous Python takes 150s.
|
|
175
|
+
|
|
176
|
+
This is 61% faster than plain synchronous Python and 19% faster than the JS baseline.
|
|
177
|
+
|
|
178
|
+
## GIL peskiness
|
|
179
|
+
|
|
180
|
+
In the current implementation we ship background calls to a `ThreadPoolExecutor`.
|
|
181
|
+
If your `@cacheable` does I/O or network work (like `requests.get`) or GIL-free CPU work (`numpy`) then you're good to go.
|
|
182
|
+
If your `@cacheable` is CPU-bound and also GIL-constrained then no dice.
|
|
183
|
+
It's not really a difficult hurdle to implement - either by switching to `ProcessPoolExecutor` or
|
|
184
|
+
`InterpreterPoolExecutor` (but that's been added only in Python 3.14), but it's not currently an option in the code.
|
|
185
|
+
|
|
186
|
+
## Disclaimer
|
|
187
|
+
|
|
188
|
+
This is a Proof-of-Concept.
|
|
189
|
+
Further work needs to be done to adjust this for other CPython versions.
|