10000 [Build] Optimize store build control for wheel and local build by ShangmingCai · Pull Request #531 · kvcache-ai/Mooncake · GitHub
[go: up one dir, main page]
More Web Proxy on the site http://driver.im/
Skip to content

[Build] Optimize store build control for wheel and local build #531

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.

Sign up for GitHub

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

Merged
merged 3 commits into from
Jun 20, 2025
Merged

[Build] Optimize store build control for wheel and local build #531

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
3f4eac6
[Build] Optimize store build control for wheel and local build
ShangmingCai Jun 20, 2025
212cc75
fix typo
ShangmingCai Jun 20, 2025
ed7205b
fix rm
ShangmingCai Jun 20, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .gitignore
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -192,5 +192,6 @@ cmake-build

libetcd_wrapper.h

mooncake-wheel/mooncake/allocator.py
mooncake-wheel/mooncake/mooncake_master
mooncake-wheel/mooncake/transfer_engine_bench
39 changes: 24 additions & 15 deletions mooncake-integration/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,11 +10,13 @@ if("${PYTHON_SYS_PATH}" STREQUAL "")
message(FATAL_ERROR "Python path is empty! Please check the python env.")
endif()

include_directories("../mooncake-store/include")
include_directories("../mooncake-store/include/cachelib_memory_allocator")
if (WITH_STORE)
include_directories("../mooncake-store/include")
include_directories("../mooncake-store/include/cachelib_memory_allocator")

include_directories("../mooncake-store/include/cachelib_memory_allocator/include")
include_directories("../mooncake-store/include/cachelib_memory_allocator/fake_include")
include_directories("../mooncake-store/include/cachelib_memory_allocator/include")
include_directories("../mooncake-store/include/cachelib_memory_allocator/fake_include")
endif()

include_directories("/usr/include/jsoncpp")

Expand Down Expand Up @@ -45,16 +47,19 @@ if(USE_MNNVL)
)
endif()

pybind11_add_module(store ${SOURCES} ${CACHE_ALLOCATOR_SOURCES}
store/store_py.cpp
)
target_link_libraries(store PUBLIC
transfer_engine
glog::glog
gflags::gflags
mooncake_store
cachelib_memory_allocator
)
if (WITH_STORE)
pybind11_add_module(store ${SOURCES} ${CACHE_ALLOCATOR_SOURCES}
store/store_py.cpp
)
target_link_libraries(store PUBLIC
transfer_engine
glog::glog
gflags::gflags
mooncake_store
cachelib_memory_allocator
)
endif()

message("${PYTHON_SYS_PATH}")

file(WRITE ${CMAKE_CURRENT_BINARY_DIR}/${PYTHON_PACKAGE_NAME}/__init__.py
Expand Down Expand Up @@ -82,5 +87,9 @@ install(
execute_process(COMMAND chmod 766 \"${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME}/__init__.py\")
"
)
install(TARGETS store DESTINATION ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME})

if (WITH_STORE)
install(TARGETS store DESTINATION ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME})
endif()

install(TARGETS engine DESTINATION ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME})
6 changes: 3 additions & 3 deletions mooncake-integration/allocator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,12 +3,12 @@
from importlib import resources
from typing import Dict, Final, Optional

import torch
from torch import device as torch_device
from torch.cuda.memory import CUDAPluggableAllocator


class NVLinkAllocator:
_instances: Dict[torch.device, CUDAPluggableAllocator] = {}
_instances: Dict[torch_device, CUDAPluggableAllocator] = {}
_lock: Final = threading.Lock()

@classmethod
Expand Down Expand Up @@ -36,7 +36,7 @@ def _get_so_path(cls) -> str:
)

@classmethod
def get_allocator(cls, device: torch.device) -> CUDAPluggableAllocator:
def get_allocator(cls, device: torch_device) -> CUDAPluggableAllocator:
with cls._lock:
if device not in cls._instances:
so_path = cls._get_so_path()
Expand Down
31 changes: 21 additions & 10 deletions scripts/build_wheel.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,29 +15,40 @@ echo "Building wheel for Python ${PYTHON_VERSION} with output directory ${OUTPUT
# Ensure LD_LIBRARY_PATH includes /usr/local/lib
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib

echo "Cleaning wheel-build directory"
rm -rf mooncake-wheel/mooncake_transfer_engine*
rm -rf mooncake-wheel/build/
rm -f mooncake-wheel/mooncake/*.so

echo "Creating directory structure..."

# Copy engine.so to mooncake directory (will be imported by transfer module)
cp build/mooncake-integration/engine.*.so mooncake-wheel/mooncake/engine.so

# Copy engine.so to mooncake directory (will be imported by transfer module)
cp build/mooncake-integration/store.*.so mooncake-wheel/mooncake/store.so
# Copy store.so to mooncake directory
if [ -f build/mooncake-integration/store.*.so ]; then
echo "Copying store.so..."
cp build/mooncake-integration/store.*.so mooncake-wheel/mooncake/store.so
echo "Copying master binary..."
# Copy master binary
cp build/mooncake-store/src/mooncake_master mooncake-wheel/mooncake/
else
echo "Skipping store.so (not built - likely WITH_STORE is set to OFF)"
fi

# Copy nvlink-allocator.so to mooncake directory (only if it exists - CUDA builds only)
if [ -f build/mooncake-transfer-engine/nvlink-allocator/nvlink_allocator.so ]; then
echo "Copying CUDA nvlink_allocator.so ..."
echo "Copying CUDA nvlink_allocator.so..."
cp build/mooncake-transfer-engine/nvlink-allocator/nvlink_allocator.so mooncake-wheel/mooncake/nvlink_allocator.so
echo "Copying allocator libraries..."
# Copy allocator.py
cp mooncake-integration/allocator.py mooncake-wheel/mooncake/allocator.py
else
echo "Skipping nvlink_allocator.so (not built - likely ARM64 or non-CUDA build)"
fi

echo "Copying allocator libraries..."
# Copy allocator.py
cp mooncake-integration/allocator.py mooncake-wheel/mooncake/allocator.py

echo "Copying master binary and shared libraries..."
# Copy master binary and shared libraries
cp build/mooncake-store/src/mooncake_master mooncake-wheel/mooncake/
echo "Copying transfer_engine_bench..."
# Copy transfer_engine_bench
cp build/mooncake-transfer-engine/example/transfer_engine_bench mooncake-wheel/mooncake/


Expand Down
Loading
0