Fixes #4016 by using websockets with callback so user can manage freeing up of state objects

[pseudotensor]

Description

Fixes #4016 by using websockets with callback so user can manage freeing up of state objects.

Closes: #4016

🎯 PRs Should Target Issues

Tests

I'd prefer to not add specific tests or other things until there is high likelihood of being merged if tests were added.

Example use cases we could test:

• Chroma database in state, callback removes the database and garbage collects to release the state
• Torch model state is in state on CUDA device, callback moves to cpu(), clears torch cache, then garbage collects to release GPU memory

Note

• I'm not a web developer, I just did my best. This PR likely needs beautification from gradio team.
• Accessing ._data private is not nice, but no other way I saw.
• One could add option in python client to setup websocket too, but this is slightly less critical because client knows their session hash, while a wild web user does not and cannot re-use that hash easily.

[gradio-pr-bot]

🪼 branch checks and previews

•	Name	Status	URL
	Spaces	ready!	Spaces preview
	Website	ready!	Website preview
	Storybook	ready!	Storybook preview
🦄	Changes	failed!	Workflow log

---

Install Gradio from this PR

pip install https://gradio-builds.s3.amazonaws.com/ad10b4a1e27c139a7b6d3f4281898d2e91788175/gradio-4.19.1-py3-none-any.whl

Install Gradio Python Client from this PR

pip install "gradio-client @ git+https://github.com/gradio-app/gradio@ad10b4a1e27c139a7b6d3f4281898d2e91788175#subdirectory=client/python"

[pseudotensor]

@abidlabs let me know what you think. As I mentioned, probably needs some love, I'm not web dev.

[abidlabs]

Thanks @pseudotensor for this PR! My main hesitation is that this solution requires the use of websockets, which we intentionally removed in the migration from Gradio 3.x -> Gradio 4.x because websockets are not supported on some cloud deployment platforms. I would be loathe to bring them back again. I suppose we could do a similar approach using http long polling instead?

Will cc @pngwn @aliabid94 @freddyaboulton for their thoughts -- generally speaking, I do think we should come up with a way to free gr.State and other temporary files when a user has finished using a Gradio app

[pseudotensor]

Maybe after your conversion way from websockets you have some ideas for how to do the equivalent. But the code is very simple and general via the callback, so it's pretty useful. I'll use it for my purposes.

Temporary files are much less of an issue than in-memory things IMHO.

[freddyaboulton]

Thanks for the PR @pseudotensor ! I think we can accomplish the same with a new SSE stream that's open once per app when it first connects. We can detect when the incoming requests disconnects and clear the state then.

[pngwn]

I'm not convinced this is reliable. SSE/ EventSource has retries / reconnections built in, that's one of the big advantages. Clearing up state based on that could easily break things. I think we need real sessions for this.

[pngwn]

Thanks for making this PR!

We don't want to go with this particularl implementation as it brings another protocol into the codebase which we would rather avoid but we definitely want to solve this, possibly using a similar approach with SSE.

Thanks again @pseudotensor!