Jupyter Code Executor in v0.4 (alternative implementation)

[Leon0402]

Why are these changes needed?

The new v4 version has no support for jupter notebooks yet compared to the old version.

NOTE: This is an alternative implementation to #4795 which uses the nbclient library rather than implementing the protocol ourself. While nbclient is probably not made for that use case, it works quite well with only small hacks. Additionally no explicit external jupyter gateway server is used anymore (which could be removed in the other implementation as well though).
Compared to the other implementation it is much simpler and more importantly much more robust. I extensively tested both implementations in a real life scenario and the other one had problems with getting frequent timeouts and got stuck after running for a while. I debugged it quite a bit, but I had the feeling that providing a robust implementation is quite complicated and better left to experts in that domain (i.e. nblient).

Related issue number

Closes #4792

Checks

• I've included any doc changes needed for https://microsoft.github.io/autogen/. See https://microsoft.github.io/autogen/docs/Contribute#documentation to build and test documentation locally.
• I've added tests (if relevant) corresponding to the changes introduced in this PR.
• I've made sure all auto checks have passed.

[jackgerrits]

This looks great! I agree, this feels like the better route to take. Simpler, more robust and offloads the complexity of processing the notebook to a library built for the task!

[Leon0402]

@jackgerrits Updated the code, I was missing some local changes that I had done in nbclient.

But there is some memory leak somewhere. When I run my program, I can see memory grow extremely. After 40min it was multiple GB. I run 10 tasks in parallel, where each task is a dialog with up to 20 turns. While I work a lot with images, which could lead to spikes in individual dialogs, the memory consumption should not be increasing with time, but be rather constant or have (unpredictable) spikes.
Unfortunately, I don't really know if it is related to this patch or some general issue in AutoGen. I already spent significant time on debugging, but i find it hard to find the exact cause. Have you experienced any similar issue without my PR here?
I could use some help :)

Some context:

async def run_task(cfg: Config, task: TaskType):
    ...

    semaphore = asyncio.Semaphore(cfg.concurrency_limit)

    async def run_single_sample(task_runner: TaskRunner, task_sample: TaskSample):
        async with semaphore:
            await task_runner.run_agent(task_sample, cfg.output_dir / task.value)

    task_runner = TaskRunner(cfg)
    samples = [run_single_sample(task_runner, task_sample) for task_sample in sliced_samples]
    await tqdm.gather(*samples, desc=f"Task: {task.value}")


async def run_tasks(cfg: Config):
    for task in cfg.tasks:
        await run_task(cfg, task)

This is my code task_runner.run_agent will then define the agents and code executor (in a async with block). So I would generally expect that each task frees up memory shortly after finishing in tqdm.gather.

[codecov]

Codecov Report

Attention: Patch coverage is 96.00000% with 4 lines in your changes missing coverage. Please review.

Project coverage is 68.44%. Comparing base (e8797a2) to head (99ff064).
Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
...t/code_executors/jupyter/_jupyter_code_executor.py	95.91%	4 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #4885      +/-   ##
==========================================
+ Coverage   68.21%   68.44%   +0.23%     
==========================================
  Files         158      160       +2     
  Lines        9960    10060     +100     
==========================================
+ Hits         6794     6886      +92     
- Misses       3166     3174       +8     

Flag	Coverage Δ
unittests	68.44% <96.00%> (+0.23%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[ekzhu]

@Leon0402 Can you show where your runtime is created? this might be due to the runtime is not removing references to created agents.

To mitigate you might want to create new instances of runtime for each task.

I think we should handle it in a separate PR.

[jackgerrits]

Haven't seen this in the main branch myself, would be curious to know what going on. I agree with Eric, sounds plausible