Skip to content
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

Merged
merged 39 commits into from
Jul 10, 2023
Merged

Migrate to Dask #244

merged 39 commits into from
Jul 10, 2023

Conversation

dlqqq
Copy link
Member

@dlqqq dlqqq commented Jun 27, 2023

Description

Replaces Ray with asyncio where possible and Dask where necessary.

Additional changes:

  • Raises Python version ceiling to 3.11
  • Bumps LangChain dependency to v0.0.220
  • Updates documentation
  • Removes more unused code from the era of model engines
    • Task API, Prompt API, TaskManager, Inserters.

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:

Task Prompt
Initialize server extension N/A
Default "Hello!"
/ask "What is Jupyter AI?"
/learn "./docs"
/generate "a notebook about asyncio"

main was profiled using #246, which adds performance profiling on main.

Task latency on main

Task Run 1 Run 2 Run 3 Avg. Latency (ms) Avg. Latency (s)
Initialize server extension 2194 2211 2218 2208 2.208
Default 2091 1129 1186 1469 1.469
/ask 4663 14452 17830 12315 12.315
/learn 13459 12856 10248 12188 12.188
/generate 269631 76938 87915 144828 144.828

Task latency on dask (current branch)

Task Run 1 Run 2 Run 3 Avg. Latency (ms) Avg. Latency (s)
Initialize server extension 22 21 21 21 0.021
Default 2125 1117 617 1286 1.286
/ask 4746 13394 10390 9510 9.51
/learn 2446 907 1460 1604 1.604
/generate 28099 55155 51436 44897 44.897

Summary of performance profiling

Task Avg. Latency (ms, main) Avg. Latency (ms, dask) Speedup Factor
Initialize server extension 2208 21 105.1
Default 1469 1286 1.1
/ask 12315 9510 1.3
/learn 12188 1604 7.6
/generate 144828 44897 3.2
  • Server initialization time was greatly decreased by a very large constant term because this branch no longer forks/spawns processes on init.
  • /learn and /generate enjoy a generous speedup thanks to rewriting actors handling network requests serially to chat handlers handling network requests concurrently.
  • The default and /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:

Screenshot 2023-06-29 at 11 01 57 AM
Mermaid code
graph TD;
    RCH["ChatHandler (root chat handler)"];
    CH[GlobalConfigHandler];
    R[RouterActor];
    A[AskActor];
    L[LearnActor];
    G[GenerateActor];
    D[DefaultActor];
    C[ConfigActor];
    M[MemoryActor];
    CP[ChatProviderActor];
    EP[EmbeddingsProviderActor];
    P[ProvidersActor];
    
    RCH --> R;
    R --> A & L & G & D & C;
    CH --> C;
    D --> M;
    A --> CP;
    L --> EP;
    CP & EP --> P;
    C --> CP & EP;

Here is how this PR migrates each actor:

  • Router actor is now implemented into the _route() method on the root chat handler, i.e. the one defined in handlers.py
  • Ask, Learn, Generate, Default are now implemented as chat handlers under chat_handlers/.
    • Base actor now is BaseChatHandler, defined in chat_handlers/base.py
  • Config actor is now implemented in ConfigManager, defined in config_manager.py
  • Chat provider, Embeddings provider, and Provider actors were all merged into ConfigManager
    • This transformation is quite natural; retrieving a provider is similar to "deserializing" a field from the configuration, which implies that this type of operation should be a method on whatever class is wrapping the configuration.
  • Memory actor is now merged into DefaultChatHandler, which is fine since nobody else was using it anyways.

The call flow now looks like this:

Screenshot 2023-06-29 at 11 02 44 AM
Mermaid code
graph TD;
    RCH[RootChatHandler];
    CH[GlobalConfigHandler];
    A[AskChatHandler];
    L[LearnChatHandler];
    G[GenerateChatHandler];
    D[DefaultChatHandler];
    C[ClearChatHandler];
    CM[ConfigManager];

    
    RCH --> A & L & G & D & C;
    CH --> CM;
    A & L & G & D --> CM;

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.

@JasonWeill
Copy link
Collaborator

By adding Python 3.11 support, this would fix #235.

@dlqqq dlqqq added enhancement New feature or request and removed enhancement New feature or request labels Jun 29, 2023
@3coins
Copy link
Collaborator

3coins commented Jun 30, 2023

@dlqqq
Thanks for tackling this. This is amazing work! 🚀

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 ().

@dlqqq
Copy link
Member Author

dlqqq commented Jun 30, 2023

@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.

@3coins
Copy link
Collaborator

3coins commented Jun 30, 2023

@dlqqq
Thanks for making that update.
I am still getting some errors on using /learn, when using Cohere large as the embedding model.

[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

@3coins
Copy link
Collaborator

3coins commented Jun 30, 2023

See some other errors with /learn and OpenAI text-embedding-ada-002 as embedding model.

 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 Cohere large and then switching to OpenAI. On restarting, this seems to have resolved, but there could be an underlying issue with failing model and then switching to a working one.

@3coins
Copy link
Collaborator

3coins commented Jun 30, 2023

After switching to OpenAI as embedding provider and restarting, /learn seems to work, but the /ask command is failing with Anthropic claude-v1.2 as the language provider.

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 Cohere xlarge and AI21 as language models.

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.

@dlqqq
Copy link
Member Author

dlqqq commented Jul 5, 2023

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 /learn and /ask.

I'll fix the LM provider bugs after standup at 11.

@dlqqq
Copy link
Member Author

dlqqq commented Jul 5, 2023

@3coins I've implemented async def _acall() for AI21 and Cohere. The issue w/ Anthropic that you've been seeing is an upstream issue with the Anthropic client, which was recently re-written with lots of breaking changes. I've pinned the version of the anthropic package we are pulling to avoid this until LangChain is compatible with the latest version of anthropic.

Tracking issue: langchain-ai/langchain#6883

@dlqqq
Copy link
Member Author

dlqqq commented Jul 5, 2023

I've rebased this branch onto main to include the fix in #249. I noticed a subtle bug in #249: self.llm_params is actually never written to; you can verify this by searching the codebase for the token "self.llm_params". This means that the LLM chain is always recreated by this conditional block:

        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 lm_provider_params is updated to None, get_llm_chain() does not return a new chain to reflect that update. I've added a one-liner that writes to self.llm_params to fix this issue: f35e75c

Copy link
Collaborator

@JasonWeill JasonWeill left a 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!

@dlqqq dlqqq merged commit d324673 into jupyterlab:main Jul 10, 2023
@dlqqq dlqqq deleted the dask branch July 27, 2023 17:44
dbelgrod pushed a commit to dbelgrod/jupyter-ai that referenced this pull request Jun 10, 2024
* 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
Marchlak pushed a commit to Marchlak/jupyter-ai that referenced this pull request Oct 28, 2024
* 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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Migrate to Dask Python 3.11 not supported
3 participants