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

All modified and coverable lines are covered by tests ✅

Project coverage is 72.95%. Comparing base (9570e82) to head (8397d22).
Report is 2 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #4885      +/-   ##
==========================================
+ Coverage   68.53%   72.95%   +4.42%     
==========================================
  Files         156      115      -41     
  Lines       10140     6784    -3356     
==========================================
- Hits         6949     4949    -2000     
+ Misses       3191     1835    -1356     

Flag	Coverage Δ
unittests	72.95% <ø> (+4.42%)	⬆️

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

[Leon0402]

@ekzhu @jackgerrits Now v4 is released, do you think we could merge this in?

[ekzhu]

Thanks. Yes let's make this part of 0.4.2. Let me review it.

[ekzhu]

@Leon0402 I pushed the changes for API docs, you can take a look, and use it as an example for future PRs: 8397d22

[Leon0402]

@ekzhu Added another example using CodeExeuctorAgent. Also used a context manager in your second example. While it is possible to use start / stop instead of a context manager, uses should only use it if they know what you they are doing. In the example it would leak resources in the case of exceptions. Thus, only do it like this in a try-catch-finally block or as part of your own context manager. For these simple examples, we should use the recommended and safe way: With Block.

Would it be possible to run the workflows automatically without maintainer approval? It is somewhat annoying that I can only see failing checks with a delay.

[ekzhu]

Would it be possible to run the workflows automatically without maintainer approval? It is somewhat annoying that I can only see failing checks with a delay.

We will work on it :)

Meanwhile you can run the checks locally by poe checks from the python directory. For more info see the developer guide under python directory.

For doc build check use poe docs-check under python/packages/autogen-core