[WIP][EP] Make uccl.ep a self-contained Python subpackage (.so + Python code) by zhenhuang12 · Pull Request #888 · uccl-project/uccl

zhenhuang12 · 2026-04-17T04:09:42Z

Description

Turns uccl.ep into a self-contained Python subpackage: the native extension (ep_cpp.abi3.so, renamed from the old top-level ep.abi3.so so it no longer collides with the uccl.ep package name) now lives alongside the Python helpers (Buffer, EventOverlap, initialize_uccl, destroy_uccl, …) under uccl/ep/. The Python source of truth moves to ep/python/uccl_ep/ and is exposed via package_dir={"uccl.ep": "ep/python/uccl_ep"}. ep/bench/{buffer,utils}.py and ep/deep_ep_wrapper/deep_ep/ are reduced to thin
re-export shims over uccl.ep, so existing bench scripts and deep_ep imports keep working. After this change pip install uccl is enough to import uccl.ep and use the full Python API.

Important

deep_ep_wrapper no longer needs to be reinstalled when the uccl.ep native library changes.
Previously deep_ep_wrapper shipped its own full copy of buffer.py / utils.py that called into the native extension directly, so any change to the uccl.ep dynamic-library API (symbol rename, signature change, new helper, …) required rebuilding and reinstalling deep_ep_wrapper to stay in sync. With this PR, deep_ep_wrapper contains no implementation of its own — it simply re-exports from uccl.ep at import time. Upgrading uccl alone is sufficient; the existing deep_ep_wrapper install picks up the new API automatically as long as the public names it re-exports remain stable.

Not a behavior change

Only packaging / module layout is changing. The native C++/CUDA code, the Python Buffer class,
initialize_uccl / destroy_uccl, benchmark scripts, and CI workflows all keep the same public semantics.
Fixes # (issue)

Type of Change

Bug fix
New feature
Documentation update

How Has This Been Tested?

Unit tests
Integration tests
Manual testing

Checklist

I have run format.sh to follow the style guidelines.
I have run build.sh to verify compilation.
I have removed redundant variables and comments.
I have updated the documentation.
I have added tests.

YangZhou1997 · 2026-04-30T16:32:43Z

There seems to be some conflict, cc @zhenhuang12

+            else:
+                try:
+                    path.unlink()
+                except FileNotFoundError:


…hon code Previously the uccl wheel only shipped the ep*.so native extension, while the Python-level functionality (Buffer, EventOverlap, initialize_uccl, etc.) lived in ep/deep_ep_wrapper and ep/bench/ — outside the uccl package. Users had to install deep_ep_wrapper separately to get the Python API. This commit restructures ep into a proper uccl.ep subpackage: 1. Rename C++ NB_MODULE from 'ep' to '_ep_native' so it becomes uccl.ep._ep_native (avoids shadowing the Python package). 2. Create uccl/ep/ with: - __init__.py: re-exports all native symbols + Python wrappers - buffer.py: the full-featured Buffer class (from ep/bench/buffer.py) - utils.py: EventOverlap, initialize_uccl, destroy_uccl, etc. 3. Update build system: - ep/setup.py: builds _ep_native, installs to site-packages/uccl/ep/ - ep/Makefile: same naming/path changes - build_inner.sh: copies .so to uccl/ep/ instead of uccl/ - Root setup.py: declares uccl.ep package_data - MANIFEST.in: includes uccl/ep/ files 4. Update deep_ep_wrapper: now purely re-exports from uccl.ep (no own logic). 5. Update ep/bench/ test scripts to use new import paths. After this change, 'pip install uccl' provides the complete EP API: from uccl.ep import Buffer, Config, EventHandle, initialize_uccl No separate deep_ep_wrapper installation is needed. Co-authored-by: zhenhuang12 <zhenhuang12@users.noreply.github.com>

ep/bench/buffer.py and ep/bench/utils.py were full copies of the implementations now canonical in uccl/ep/buffer.py and uccl/ep/utils.py. Having two copies means maintaining the same code in two places. Replace them with thin shim modules that re-export everything from uccl.ep.{buffer,utils}. The bench test scripts (test_low_latency.py, test_intranode.py, etc.) continue to do 'from buffer import Buffer' and 'from utils import ...' which resolves through the shims. This establishes uccl/ep/ as the single source of truth. Co-authored-by: zhenhuang12 <zhenhuang12@users.noreply.github.com>

…of truth) The previous approach duplicated buffer.py and utils.py in both uccl/ep/ and ep/bench/. This commit eliminates the duplication by establishing a clear convention: - Source of truth: ep/python/uccl_ep/{__init__,buffer,utils}.py (lives alongside the ep module's C++/CUDA sources) - Installed as: uccl.ep (via setup.py package_dir mapping) How it works: 1. Root setup.py uses package_dir={'uccl.ep': 'ep/python/uccl_ep'} so setuptools reads .py files directly from the ep module tree. 2. build_inner.sh copies the compiled _ep_native*.so into ep/python/uccl_ep/ before running 'python -m build', so both .py and .so end up in the wheel. 3. ep/setup.py CustomInstall copies both .so and .py files to site-packages/uccl/ep/ for direct-install workflows. 4. ep/Makefile install target does the same for make-based workflows. 5. ep/bench/{buffer,utils}.py remain thin shims that re-export from uccl.ep, so bench scripts work unchanged. The uccl/ep/ directory is no longer checked into the repository; it only exists in the installed site-packages. Co-authored-by: zhenhuang12 <zhenhuang12@users.noreply.github.com>

_ep_native -> ep_cpp

zhenhuang12 · 2026-05-06T01:53:28Z

There seems to be some conflict, cc @zhenhuang12

@YangZhou1997 Thanks! I'm continuing to support this feature—the build changes are fairly large and I'd appreciate a review. I've moved the compilation of build_inner into the Makefile, while leaving the packaging scripts unchanged. Now we can install uccl directly via pip install . --no-build-isolation, and then use the ep module with from uccl import ep. The ep module is a complete Python package, and it has been decoupled from the deep_ep_wrapper dependency.

cc @MaoZiming

YangZhou1997 · 2026-05-06T22:44:02Z

@zhenhuang12 is it possible to avoid using a Makefile, which is usually less maintainable than a shell script like build_inner.sh?

zhenhuang12 · 2026-05-07T03:11:48Z

@zhenhuang12 is it possible to avoid using a Makefile, which is usually less maintainable than a shell script like build_inner.sh?

@YangZhou1997 I think we could replace the Makefile with a build.sh script, which is essentially equivalent to writing a ShellExtension to call build.sh in setup.py. Currently build_inner.sh covers both build and packaging steps, but we can separate out the build part from it.

Do you think using build.sh instead of Makefile is a better approach? I'd like to hear your opinion.

zhenhuang12 · 2026-05-07T03:16:28Z

@YangZhou1997 If you think the current PR change is too large, I can revert it and only keep the "ep python self-containerd" part. Support for ShellExtension or MakeExtension can then be introduced in a later PR.

YangZhou1997 · 2026-05-07T03:27:22Z

Hi @zhenhuang12, the PR size is okay for me. I think if we can use a ShellExtension in the setup.py, that would be better, so that we can keep all building functions in the existing build.sh and build_inner.sh (and get rid of Makefile in the top level)

github-code-quality Bot found potential problems Apr 17, 2026

View reviewed changes

zhenhuang12 added the WIP Work In Progress label Apr 17, 2026

github-code-quality Bot found potential problems Apr 30, 2026

View reviewed changes

Comment thread ep/deep_ep_wrapper/deep_ep/__init__.py Fixed

Comment thread ep/deep_ep_wrapper/deep_ep/__init__.py Fixed

Comment thread setup.py Fixed

zhenhuang12 requested a review from YangZhou1997 April 30, 2026 12:53

github-code-quality Bot found potential problems Apr 30, 2026

View reviewed changes

zhenhuang12 requested a review from MaoZiming April 30, 2026 12:53

zhenhuang12 force-pushed the cursor/ep-self-contained-package-59eb branch from 7fa4c37 to 217c5fd Compare April 30, 2026 12:55

github-code-quality Bot found potential problems Apr 30, 2026

View reviewed changes

Comment thread ep/python/uccl_ep/utils.py Fixed

zhenhuang12 changed the title ~~[EP] Make uccl.ep a self-contained Python subpackage (.so + Python code)~~ [WIP][EP] Make uccl.ep a self-contained Python subpackage (.so + Python code) Apr 30, 2026

YangZhou1997 added the run-benchmark label Apr 30, 2026

zhenhuang12 force-pushed the cursor/ep-self-contained-package-59eb branch from f056922 to 9f26d95 Compare May 5, 2026 01:09

github-code-quality Bot found potential problems May 5, 2026

View reviewed changes

Comment thread ep/python/uccl_ep/utils.py Fixed

Comment thread ep/bench/utils.py Fixed

Comment thread ep/python/uccl_ep/utils.py Fixed

Comment thread ep/python/uccl_ep/utils.py Fixed

Comment thread ep/python/uccl_ep/utils.py Fixed

Comment thread ep/python/uccl_ep/utils.py Fixed

github-code-quality Bot found potential problems May 6, 2026

View reviewed changes

Comment thread setup.py

else:

try:

path.unlink()

except FileNotFoundError:

cursoragent and others added 12 commits May 6, 2026 01:37

refactor: rename native EP extension

46fc398

_ep_native -> ep_cpp

add make extension

843a59c

refactor

0c88586

add doc

2807f96

rebase

b755942

nits

bf86d1e

nits

29c42be

nits

afe786d

refactor build_inner: build uccl by makefile

ccd6ee2

solve conflict

9bae94c

zhenhuang12 force-pushed the cursor/ep-self-contained-package-59eb branch from 9a87824 to 9bae94c Compare May 6, 2026 02:10

nits

f6ddf4d

zhenhuang12 and others added 3 commits May 6, 2026 03:48

nits

8a197fd

add back

f348162

fixing wheel missing issues in docker containers

5dee1a2

YangZhou1997 added 2 commits May 6, 2026 23:00

fixing gb10 test

a5e86c8

fixing build on l4

1085b99

Conversation

zhenhuang12 commented Apr 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Not a behavior change

Type of Change

How Has This Been Tested?

Checklist

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

YangZhou1997 commented Apr 30, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

zhenhuang12 commented May 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

YangZhou1997 commented May 6, 2026

Uh oh!

zhenhuang12 commented May 7, 2026

Uh oh!

zhenhuang12 commented May 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

YangZhou1997 commented May 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

zhenhuang12 commented Apr 17, 2026 •

edited

Loading

zhenhuang12 commented May 6, 2026 •

edited

Loading

zhenhuang12 commented May 7, 2026 •

edited

Loading

YangZhou1997 commented May 7, 2026 •

edited

Loading