-
-
Notifications
You must be signed in to change notification settings - Fork 332
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Migrate to Dask #244
Migrate to Dask #244
Conversation
By adding Python 3.11 support, this would fix #235. |
@dlqqq For Python 3.9, this is failing installation. It might be related to extending the EmbeddingsProviders from the Langchain providers. Error during installation@jupyter-ai/core: An error occurred.
@jupyter-ai/core: ModuleNotFoundError: There is no labextension at .. Errors encountered: [TypeError("the 'package' argument is required to perform a relative import for '.'"), TypeError('Tuple[t0, t1, ...]: each t must be a type. Got ().'), TypeError('Tuple[t0, t1, ...]: each t must be a type. Got ().'), TypeError('Tuple[t0, t1, ...]: each t must be a type. Got ().'), TypeError('Tuple[t0, t1, ...]: each t must be a type. Got ().'), TypeError('Tuple[t0, t1, ...]: each t must be a type. Got ().'), TypeError('Tuple[t0, t1, ...]: each t must be a type. Got ().'), TypeError('Tuple[t0, t1, ...]: each t must be a type. Got ().'), TypeError('Tuple[t0, t1, ...]: each t must be a type. Got ().'), TypeError('Tuple[t0, t1, ...]: each t must be a type. Got ().'), TypeError('Tuple[t0, t1, ...]: each t must be a type. Got ().'), TypeError('Tuple[t0, t1, ...]: each t must be a type. Got ().'), TypeError('Tuple[t0, t1, ...]: each t must be a type. Got ().'), TypeError('Tuple[t0, t1, ...]: each t must be a type. Got ().'), TypeError('Tuple[t0, t1, ...]: each t must be a type. Got ().')]
@jupyter-ai/core: See the log file for details: /var/folders/zz/mkx8_cg176xdvqx28gqt05wr0000gr/T/jupyterlab-debug-g8lqr9_l.log
error Command failed with exit code 1.
info Visit https://yarnpkg.com/en/docs/cli/run for documentation about this command. Error on server start[W 2023-06-29 22:00:45.878 ServerApp] jupyter_ai | error adding extension (enabled: True): Tuple[t0, t1, ...]: each t must be a type. Got ().
Traceback (most recent call last):
File "/opt/anaconda3/envs/jupyter-ai-with-dask/lib/python3.9/site-packages/jupyter_server/extension/manager.py", line 319, in add_extension
extpkg = ExtensionPackage(name=extension_name, enabled=enabled)
File "/opt/anaconda3/envs/jupyter-ai-with-dask/lib/python3.9/site-packages/jupyter_server/extension/manager.py", line 183, in __init__
self._load_metadata()
File "/opt/anaconda3/envs/jupyter-ai-with-dask/lib/python3.9/site-packages/jupyter_server/extension/manager.py", line 192, in _load_metadata
self.module, self.metadata = get_metadata(name, logger=self.log)
File "/opt/anaconda3/envs/jupyter-ai-with-dask/lib/python3.9/site-packages/jupyter_server/extension/utils.py", line 72, in get_metadata
module = importlib.import_module(package_name)
File "/opt/anaconda3/envs/jupyter-ai-with-dask/lib/python3.9/importlib/__init__.py", line 127, in import_module
return _bootstrap._gcd_import(name[level:], package, level)
File "<frozen importlib._bootstrap>", line 1030, in _gcd_import
File "<frozen importlib._bootstrap>", line 1007, in _find_and_load
File "<frozen importlib._bootstrap>", line 986, in _find_and_load_unlocked
File "<frozen importlib._bootstrap>", line 680, in _load_unlocked
File "<frozen importlib._bootstrap_external>", line 850, in exec_module
File "<frozen importlib._bootstrap>", line 228, in _call_with_frames_removed
File "/Users/pijain/projects/jai-with-dask/jupyter-ai/packages/jupyter-ai/jupyter_ai/__init__.py", line 2, in <module>
from jupyter_ai_magics import load_ipython_extension, unload_ipython_extension
File "/Users/pijain/projects/jai-with-dask/jupyter-ai/packages/jupyter-ai-magics/jupyter_ai_magics/__init__.py", line 4, in <module>
from .embedding_providers import (
File "/Users/pijain/projects/jai-with-dask/jupyter-ai/packages/jupyter-ai-magics/jupyter_ai_magics/embedding_providers.py", line 62, in <module>
class OpenAIEmbeddingsProvider(BaseEmbeddingsProvider, OpenAIEmbeddings):
File "pydantic/main.py", line 138, in pydantic.main.ModelMetaclass.__new__
File "pydantic/utils.py", line 693, in pydantic.utils.smart_deepcopy
File "/opt/anaconda3/envs/jupyter-ai-with-dask/lib/python3.9/copy.py", line 146, in deepcopy
y = copier(x, memo)
File "/opt/anaconda3/envs/jupyter-ai-with-dask/lib/python3.9/copy.py", line 230, in _deepcopy_dict
y[deepcopy(key, memo)] = deepcopy(value, memo)
File "/opt/anaconda3/envs/jupyter-ai-with-dask/lib/python3.9/copy.py", line 172, in deepcopy
y = _reconstruct(x, memo, *rv)
File "/opt/anaconda3/envs/jupyter-ai-with-dask/lib/python3.9/copy.py", line 270, in _reconstruct
state = deepcopy(state, memo)
File "/opt/anaconda3/envs/jupyter-ai-with-dask/lib/python3.9/copy.py", line 146, in deepcopy
y = copier(x, memo)
File "/opt/anaconda3/envs/jupyter-ai-with-dask/lib/python3.9/copy.py", line 210, in _deepcopy_tuple
y = [deepcopy(a, memo) for a in x]
File "/opt/anaconda3/envs/jupyter-ai-with-dask/lib/python3.9/copy.py", line 210, in <listcomp>
y = [deepcopy(a, memo) for a in x]
File "/opt/anaconda3/envs/jupyter-ai-with-dask/lib/python3.9/copy.py", line 146, in deepcopy
y = copier(x, memo)
File "/opt/anaconda3/envs/jupyter-ai-with-dask/lib/python3.9/copy.py", line 230, in _deepcopy_dict
y[deepcopy(key, memo)] = deepcopy(value, memo)
File "/opt/anaconda3/envs/jupyter-ai-with-dask/lib/python3.9/copy.py", line 172, in deepcopy
y = _reconstruct(x, memo, *rv)
File "/opt/anaconda3/envs/jupyter-ai-with-dask/lib/python3.9/copy.py", line 264, in _reconstruct
y = func(*args)
File "/opt/anaconda3/envs/jupyter-ai-with-dask/lib/python3.9/copy.py", line 263, in <genexpr>
args = (deepcopy(arg, memo) for arg in args)
File "/opt/anaconda3/envs/jupyter-ai-with-dask/lib/python3.9/copy.py", line 146, in deepcopy
y = copier(x, memo)
File "/opt/anaconda3/envs/jupyter-ai-with-dask/lib/python3.9/copy.py", line 210, in _deepcopy_tuple
y = [deepcopy(a, memo) for a in x]
File "/opt/anaconda3/envs/jupyter-ai-with-dask/lib/python3.9/copy.py", line 210, in <listcomp>
y = [deepcopy(a, memo) for a in x]
File "/opt/anaconda3/envs/jupyter-ai-with-dask/lib/python3.9/copy.py", line 172, in deepcopy
y = _reconstruct(x, memo, *rv)
File "/opt/anaconda3/envs/jupyter-ai-with-dask/lib/python3.9/copy.py", line 264, in _reconstruct
y = func(*args)
File "/opt/anaconda3/envs/jupyter-ai-with-dask/lib/python3.9/typing.py", line 277, in inner
return func(*args, **kwds)
File "/opt/anaconda3/envs/jupyter-ai-with-dask/lib/python3.9/typing.py", line 920, in __getitem__
params = tuple(_type_check(p, msg) for p in params)
File "/opt/anaconda3/envs/jupyter-ai-with-dask/lib/python3.9/typing.py", line 920, in <genexpr>
params = tuple(_type_check(p, msg) for p in params)
File "/opt/anaconda3/envs/jupyter-ai-with-dask/lib/python3.9/typing.py", line 166, in _type_check
raise TypeError(f"{msg} Got {arg!r:.100}.")
TypeError: Tuple[t0, t1, ...]: each t must be a type. Got (). |
@3coins Thanks for the callout. I traced the bug to langchain-ai/langchain#4500, which was fixed in v0.0.172. I've upgraded our LangChain dependency to the latest version, v0.0.220, which seems to fix this issue. |
@dlqqq [I 2023-06-30 12:21:29.756 ServerApp] /learn chat handler resolved in 463 ms.
2023-06-30 12:22:48,648 - distributed.protocol.pickle - ERROR - Failed to serialize <ToPickle: HighLevelGraph with 26 layers.
<dask.highlevelgraph.HighLevelGraph object at 0x12ee9ba30>
0. embed_chunk-efaf5a17-e05e-4a74-a842-36e974a59d02
1. embed_chunk-a15178a2-d9b7-4f1b-99a3-a65d97904d87
2. embed_chunk-9be8ad94-d0b6-4e2e-915e-61ec291983db
3. embed_chunk-21d705eb-d52c-4877-8f44-809d1aa18b75
4. embed_chunk-bc883d85-fe96-4520-a717-9faa43c871dd
5. embed_chunk-97b773ee-df92-4be7-9610-c46e6efcbc29
6. embed_chunk-bbb96722-e317-4af4-ba36-7bc62c44828b
7. embed_chunk-76558ec6-5872-4c9c-808b-9782f7aaebd1
8. embed_chunk-4155863f-5666-41d5-a2dc-fdb81b92aff2
9. embed_chunk-7f43d3c6-e78c-426c-b677-f16c096f4374
10. embed_chunk-d163a876-6704-4fb6-acbf-a6f154f41813
11. embed_chunk-b169bd11-c947-41d7-95d3-5cc862bc3054
12. embed_chunk-8de47bc8-10b0-45a9-96c5-8b263a1a0eee
13. embed_chunk-4b3de3df-a8ba-408e-ba93-808a4a018348
14. embed_chunk-693821cd-d7e5-4a30-a1d2-77b8c343d312
15. embed_chunk-eed85cbb-d211-4c52-ad24-cac3817734d1
16. embed_chunk-f937bbd2-d622-445f-bf2c-f016d1c721ab
17. embed_chunk-43950423-a52c-4723-8bef-4eb008d0442f
18. embed_chunk-fe3e544e-4207-4053-b3aa-57da85a2058a
19. embed_chunk-876eb777-98a9-41fb-8982-ebf8e57d9c05
20. embed_chunk-04f3a19d-7b53-47e9-a6ea-04a7974d067e
21. embed_chunk-d614aad0-4d92-4043-bdf1-e5dbc2e6ed0f
22. embed_chunk-bcc6ce5c-a78c-4397-a829-6298ae97d1e3
23. embed_chunk-950e51bc-e550-4fb8-93df-9b886d1c428d
24. join-864feb23-3415-4020-a1d3-343288100857
25. 1c91f8b8611bfa3c83d86dbbc001eb31
>.
Traceback (most recent call last):
File "/opt/anaconda3/envs/jupyter-ai-with-dask/lib/python3.9/site-packages/distributed/protocol/pickle.py", line 63, in dumps
result = pickle.dumps(x, **dump_kwargs)
TypeError: cannot pickle '_queue.SimpleQueue' object
During handling of the above exception, another exception occurred:
Traceback (most recent call last):
File "/opt/anaconda3/envs/jupyter-ai-with-dask/lib/python3.9/site-packages/distributed/protocol/pickle.py", line 68, in dumps
pickler.dump(x)
TypeError: cannot pickle '_queue.SimpleQueue' object
During handling of the above exception, another exception occurred:
Traceback (most recent call last):
File "/opt/anaconda3/envs/jupyter-ai-with-dask/lib/python3.9/site-packages/distributed/protocol/pickle.py", line 81, in dumps
result = cloudpickle.dumps(x, **dump_kwargs)
File "/opt/anaconda3/envs/jupyter-ai-with-dask/lib/python3.9/site-packages/cloudpickle/cloudpickle_fast.py", line 73, in dumps
cp.dump(obj)
File "/opt/anaconda3/envs/jupyter-ai-with-dask/lib/python3.9/site-packages/cloudpickle/cloudpickle_fast.py", line 632, in dump
return Pickler.dump(self, obj)
TypeError: cannot pickle '_queue.SimpleQueue' object |
See some other errors with File "/Users/pijain/projects/jai-with-dask/jupyter-ai/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py", line 166, in delete
self.create()
File "/Users/pijain/projects/jai-with-dask/jupyter-ai/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py", line 187, in create
embeddings = self.get_embedding_model()
File "/Users/pijain/projects/jai-with-dask/jupyter-ai/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py", line 254, in get_embedding_model
self.delete_and_relearn()
File "/Users/pijain/projects/jai-with-dask/jupyter-ai/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py", line 145, in delete_and_relearn
self.delete()
File "/Users/pijain/projects/jai-with-dask/jupyter-ai/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py", line 166, in delete
self.create()
File "/Users/pijain/projects/jai-with-dask/jupyter-ai/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py", line 187, in create
embeddings = self.get_embedding_model()
File "/Users/pijain/projects/jai-with-dask/jupyter-ai/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py", line 254, in get_embedding_model
self.delete_and_relearn()
File "/Users/pijain/projects/jai-with-dask/jupyter-ai/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py", line 145, in delete_and_relearn
self.delete()
File "/Users/pijain/projects/jai-with-dask/jupyter-ai/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py", line 166, in delete
self.create()
File "/Users/pijain/projects/jai-with-dask/jupyter-ai/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py", line 187, in create
embeddings = self.get_embedding_model()
File "/Users/pijain/projects/jai-with-dask/jupyter-ai/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py", line 249, in get_embedding_model
self.embeddings = em_provider(**em_provider_params)
File "/Users/pijain/projects/jai-with-dask/jupyter-ai/packages/jupyter-ai-magics/jupyter_ai_magics/embedding_providers.py", line 59, in __init__
super().__init__(*args, **kwargs, **model_kwargs)
File "pydantic/main.py", line 339, in pydantic.main.BaseModel.__init__
File "pydantic/main.py", line 1076, in pydantic.main.validate_model
File "pydantic/fields.py", line 884, in pydantic.fields.ModelField.validate
File "pydantic/fields.py", line 1101, in pydantic.fields.ModelField._validate_singleton
File "pydantic/fields.py", line 1157, in pydantic.fields.ModelField._apply_validators
File "pydantic/class_validators.py", line 337, in pydantic.class_validators._generic_validator_basic.lambda13
RecursionError: maximum recursion depth exceeded while calling a Python object Seems like this is a remnant of an error that happens after using |
After switching to Traceback (most recent call last):
File "/Users/pijain/projects/jai-with-dask/jupyter-ai/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py", line 38, in process_message
await self._process_message(message)
File "/Users/pijain/projects/jai-with-dask/jupyter-ai/packages/jupyter-ai/jupyter_ai/chat_handlers/ask.py", line 44, in _process_message
self.get_llm_chain()
File "/Users/pijain/projects/jai-with-dask/jupyter-ai/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py", line 83, in get_llm_chain
self.create_llm_chain(lm_provider, lm_provider_params)
File "/Users/pijain/projects/jai-with-dask/jupyter-ai/packages/jupyter-ai/jupyter_ai/chat_handlers/ask.py", line 29, in create_llm_chain
self.llm = provider(**provider_params)
File "/Users/pijain/projects/jai-with-dask/jupyter-ai/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py", line 121, in __init__
super().__init__(*args, **kwargs, **model_kwargs)
File "/opt/anaconda3/envs/jupyter-ai-with-dask/lib/python3.9/site-packages/langchain/load/serializable.py", line 74, in __init__
super().__init__(**kwargs)
File "pydantic/main.py", line 341, in pydantic.main.BaseModel.__init__
pydantic.error_wrappers.ValidationError: 1 validation error for AnthropicProvider
__root__
__init__() got an unexpected keyword argument 'api_url' (type=type_error) Here is the error with File "/opt/anaconda3/envs/jupyter-ai-with-dask/lib/python3.9/site-packages/langchain/llms/base.py", line 325, in agenerate
output = await self._agenerate_helper(
File "/opt/anaconda3/envs/jupyter-ai-with-dask/lib/python3.9/site-packages/langchain/llms/base.py", line 275, in _agenerate_helper
raise e
File "/opt/anaconda3/envs/jupyter-ai-with-dask/lib/python3.9/site-packages/langchain/llms/base.py", line 262, in _agenerate_helper
await self._agenerate(
File "/opt/anaconda3/envs/jupyter-ai-with-dask/lib/python3.9/site-packages/langchain/llms/base.py", line 544, in _agenerate
await self._acall(prompt, stop=stop, run_manager=run_manager, **kwargs)
File "/opt/anaconda3/envs/jupyter-ai-with-dask/lib/python3.9/site-packages/langchain/llms/base.py", line 510, in _acall
raise NotImplementedError("Async generation not implemented for this LLM.")
NotImplementedError: Async generation not implemented for this LLM. |
I've fixed the bugs relating to embedding models. You can now use the Cohere embedding provider and can freely switch between embedding providers. Jupyter AI also automatically relearns indexed directories on subsequent calls to I'll fix the LM provider bugs after standup at 11. |
@3coins I've implemented Tracking issue: langchain-ai/langchain#6883 |
I've rebased this branch onto elif self.llm_params != lm_provider_params:
self.log.info("Chat model params changed, updating the llm chain.")
self.create_llm_chain(lm_provider, lm_provider_params) This also means that if |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Tested this successfully. Thanks!
* use dask to parallelize document learning * prefer threads * remove memory actor * remove providers actor * remove chat provider and embeddings actors * remove config actor * remove router actor * move embedding logic from base to learn actor * move remaining actors to chat_handlers * remove references to ray and grpcio * implement async handlers and concurrency * precommit * remove compute_delayed() * clear chat history in clear handler * log chat handler latency * delete useless comment * make /generate async and concurrent * raise Python version ceiling to 3.11 * support Dask on Python 3.11 * pre-commit fixes * create task per message in root chat handler, fix concurrency * log server extension init time * remove unused code from the era of model engines * pre-commit fixes * bump langchain to 0.0.220 * fix recursion error when switching embedding providers * fix pickle issue and relearning logic * relearn on /ask if embedding model was changed * pre-commit * implement _acall for AI21 and Cohere providers * pin anthropic to 0.2.10 * pre-commit * write to self.llm_params after creating chain * fix LearnChatHandler to not throw on empty config * update example notebooks * return if no embedding provider selected * make SM Endpoints provider async * implement async for HF Hub provider * strip quotation marks from generated notebook titles
* use dask to parallelize document learning * prefer threads * remove memory actor * remove providers actor * remove chat provider and embeddings actors * remove config actor * remove router actor * move embedding logic from base to learn actor * move remaining actors to chat_handlers * remove references to ray and grpcio * implement async handlers and concurrency * precommit * remove compute_delayed() * clear chat history in clear handler * log chat handler latency * delete useless comment * make /generate async and concurrent * raise Python version ceiling to 3.11 * support Dask on Python 3.11 * pre-commit fixes * create task per message in root chat handler, fix concurrency * log server extension init time * remove unused code from the era of model engines * pre-commit fixes * bump langchain to 0.0.220 * fix recursion error when switching embedding providers * fix pickle issue and relearning logic * relearn on /ask if embedding model was changed * pre-commit * implement _acall for AI21 and Cohere providers * pin anthropic to 0.2.10 * pre-commit * write to self.llm_params after creating chain * fix LearnChatHandler to not throw on empty config * update example notebooks * return if no embedding provider selected * make SM Endpoints provider async * implement async for HF Hub provider * strip quotation marks from generated notebook titles
Description
Replaces Ray with asyncio where possible and Dask where necessary.
Additional changes:
Issues closed:
Demo
This demonstrates how multiprocessing is not necessary to implement concurrency and non-blocking messages.
Screen.Recording.2023-06-27.at.10.54.00.AM.mov
Performance profiling
All performance profiling was done on a M1 Macbook Pro with 32 GB unified memory. Each task was performed with the following prompts:
main
was profiled using #246, which adds performance profiling onmain
.Task latency on
main
Task latency on
dask
(current branch)Summary of performance profiling
main
)dask
)/learn
and/generate
enjoy a generous speedup thanks to rewriting actors handling network requests serially to chat handlers handling network requests concurrently./ask
latency were decreased only by an apparently constant amount of time, as the performance gain is due exclusively to removing IPC overhead.Developer guidance
Where did the actors go?
Previously, the actors had a complex call flow:
Mermaid code
Here is how this PR migrates each actor:
_route()
method on the root chat handler, i.e. the one defined inhandlers.py
chat_handlers/
.BaseChatHandler
, defined inchat_handlers/base.py
ConfigManager
, defined inconfig_manager.py
ConfigManager
DefaultChatHandler
, which is fine since nobody else was using it anyways.The call flow now looks like this:
Mermaid code
Why does this work?
This works because it's not necessary to spawn multiple processes to handle web server concurrency. Tornado itself is a single-threaded web server. Async coroutines handle concurrency very well thanks to how fast the context switches are relative to using processes. We should reserve parallelism for CPU-bound/GIL-bound parallelizable tasks (as is the case for document indexing), and prefer
asyncio
everywhere else, as typically we are I/O-bound and just need concurrency.How should I review this PR?
I recommend testing all commands and use-cases, making sure they're concurrent (i.e. allow the backend to respond to a quick message after a pending message, e.g.
/generate
), and then making sure that all Python versions are known to work.