From 42e306c28118bde224f3acc7f5981d32fd59759a Mon Sep 17 00:00:00 2001 From: Aleks Knezevic Date: Sat, 22 Feb 2025 02:33:18 +0000 Subject: [PATCH 1/2] Fix venv path, so we no longer require modification to python path after install --- CMakeLists.txt | 2 +- env/activate | 2 +- setup.py | 2 +- third_party/CMakeLists.txt | 7 ++++++- tt_torch/csrc/CMakeLists.txt | 15 +++++++++++++-- 5 files changed, 22 insertions(+), 6 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 4984951a..251eb2c2 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -31,7 +31,7 @@ set(CMAKE_EXPORT_COMPILE_COMMANDS ON) set(CMAKE_CXX_STANDARD 17) set(CMAKE_POSITION_INDEPENDENT_CODE ON) -if(NOT DEFINED ENV{DONT_OVERRIDE_INSTALL_PATH}) +if(NOT DEFINED ENV{WHEEL_BUILD_CUSTOM_INSTALL_PATH}) set(CMAKE_INSTALL_PREFIX ${TTTORCH_SOURCE_DIR}/install) endif() message(STATUS "CMAKE_INSTALL_PREFIX is set to ${CMAKE_INSTALL_PREFIX}") diff --git a/env/activate b/env/activate index 55e28220..125c4379 100644 --- a/env/activate +++ b/env/activate @@ -53,6 +53,6 @@ else export TT_METAL_HOME="$(pwd)/third_party/tt-mlir/src/tt-mlir/third_party/tt-metal/src/tt-metal" fi -export PYTHONPATH="$(pwd):$(pwd)/env/venv:$(pwd)/env/venv/lib:$(pwd)/install/lib:$(pwd)/.local/toolchain/python_packages/mlir_core:${TT_METAL_HOME}:${TT_METAL_HOME}/tt_eager:${TT_METAL_BUILD_HOME}/tools/profiler/bin" +export PYTHONPATH="$(pwd):$(pwd)/install/lib:${TT_METAL_HOME}:${TT_METAL_HOME}/tt_eager:${TT_METAL_BUILD_HOME}/tools/profiler/bin" export ARCH_NAME="${ARCH_NAME:-wormhole_b0}" export TT_METAL_LOGGER_LEVEL="ERROR" diff --git a/setup.py b/setup.py index 60b12da7..95e96765 100644 --- a/setup.py +++ b/setup.py @@ -29,7 +29,7 @@ def run(self): # Compile time env vars -os.environ["DONT_OVERRIDE_INSTALL_PATH"] = "1" +os.environ["WHEEL_BUILD_CUSTOM_INSTALL_PATH"] = "1" cmake_args = [ "-GNinja", diff --git a/third_party/CMakeLists.txt b/third_party/CMakeLists.txt index fb8aed5b..d904c1f0 100644 --- a/third_party/CMakeLists.txt +++ b/third_party/CMakeLists.txt @@ -72,6 +72,11 @@ else() set_target_properties(TTMLIRRuntime PROPERTIES EXCLUDE_FROM_ALL TRUE IMPORTED_LOCATION ${TT_MLIR_RUNTIME_LIBRARY_PATH}) add_dependencies(TTMLIRRuntime tt-mlir) - install(DIRECTORY ${TTMLIR_INSTALL_PREFIX}/ DESTINATION "${CMAKE_INSTALL_PREFIX}" USE_SOURCE_PERMISSIONS) + if(DEFINED ENV{WHEEL_BUILD_CUSTOM_INSTALL_PATH}) + file(MAKE_DIRECTORY ${CMAKE_INSTALL_PREFIX}/lib/python3.11/site-packages) + install(DIRECTORY ${TTMLIR_INSTALL_PREFIX}/ DESTINATION "${CMAKE_INSTALL_PREFIX}/lib/python3.11/site-packages" USE_SOURCE_PERMISSIONS) + else() + install(DIRECTORY ${TTMLIR_INSTALL_PREFIX}/ DESTINATION "${CMAKE_INSTALL_PREFIX}" USE_SOURCE_PERMISSIONS) + endif() endif() diff --git a/tt_torch/csrc/CMakeLists.txt b/tt_torch/csrc/CMakeLists.txt index e14dc773..e5fee8ec 100644 --- a/tt_torch/csrc/CMakeLists.txt +++ b/tt_torch/csrc/CMakeLists.txt @@ -13,7 +13,12 @@ add_dependencies(TT_TORCH_MLIR ) set_target_properties(TT_TORCH_MLIR PROPERTIES COMPILE_FLAGS "-fno-rtti") -install (TARGETS TT_TORCH_MLIR LIBRARY DESTINATION ${CMAKE_INSTALL_PREFIX}/lib) +if(DEFINED ENV{WHEEL_BUILD_CUSTOM_INSTALL_PATH}) + file(MAKE_DIRECTORY ${CMAKE_INSTALL_PREFIX}/lib/python3.11/site-packages/tt_torch) + install (TARGETS TT_TORCH_MLIR LIBRARY DESTINATION ${CMAKE_INSTALL_PREFIX}/lib/python3.11/site-packages/tt_torch) +else() + install (TARGETS TT_TORCH_MLIR LIBRARY DESTINATION ${CMAKE_INSTALL_PREFIX}/lib) +endif() target_include_directories(TT_TORCH_MLIR PUBLIC ${PROJECT_SOURCE_DIR}/third_party/tt-mlir/src/tt-mlir-build/include @@ -80,4 +85,10 @@ set_target_properties(${TARGET_NAME} LIBRARY_OUTPUT_NAME ${TARGET_NAME} INSTALL_RPATH "$ORIGIN" ) -install (TARGETS ${TARGET_NAME} LIBRARY DESTINATION ${CMAKE_INSTALL_PREFIX}/lib) + +if(DEFINED ENV{WHEEL_BUILD_CUSTOM_INSTALL_PATH}) + file(MAKE_DIRECTORY ${CMAKE_INSTALL_PREFIX}/lib/python3.11/site-packages/tt_torch) + install (TARGETS ${TARGET_NAME} LIBRARY DESTINATION ${CMAKE_INSTALL_PREFIX}/lib/python3.11/site-packages/tt_torch) +else() + install (TARGETS ${TARGET_NAME} LIBRARY DESTINATION ${CMAKE_INSTALL_PREFIX}/lib) +endif() From cb3b1aa38ff24f0ed4dc8e6391b0074d06c19462 Mon Sep 17 00:00:00 2001 From: Aleks Knezevic Date: Sun, 23 Feb 2025 03:18:54 +0000 Subject: [PATCH 2/2] try moving files within the whl --- env/activate | 2 +- setup.py | 58 +++++++++++++++++++++++++++++++++++- third_party/CMakeLists.txt | 3 +- tt_torch/__init__.py | 6 ++++ tt_torch/csrc/CMakeLists.txt | 14 ++------- tt_torch/tools/__init__.py | 0 6 files changed, 68 insertions(+), 15 deletions(-) create mode 100644 tt_torch/tools/__init__.py diff --git a/env/activate b/env/activate index 125c4379..5e2a95d5 100644 --- a/env/activate +++ b/env/activate @@ -16,7 +16,7 @@ else fi export TT_TORCH_HOME="$(pwd)" -export LD_LIBRARY_PATH=$TT_TORCH_HOME/env/venv/lib/python3.11/site-packages/torch/lib:$TTMLIR_TOOLCHAIN_DIR/lib:$TT_TORCH_HOME/install/lib/:$LD_LIBRRARY_PATH +export LD_LIBRARY_PATH=$TT_TORCH_HOME/env/venv/lib/python3.11/site-packages/torch/lib:$TTMLIR_TOOLCHAIN_DIR/lib:$TT_TORCH_HOME/install/lib/:$TT_TORCH_HOME/env/venv/lib:$LD_LIBRRARY_PATH export TTMLIR_VENV_DIR="$(pwd)/env/venv" if [ -d $TTMLIR_VENV_DIR/bin ]; then diff --git a/setup.py b/setup.py index 95e96765..67bced85 100644 --- a/setup.py +++ b/setup.py @@ -3,10 +3,64 @@ # SPDX-License-Identifier: Apache-2.0 from skbuild import setup import os -import glob from skbuild.command.install_lib import install_lib from setuptools import find_namespace_packages import sys +import shutil + +import hashlib +import zipfile + +# TODO, this is quite hacky +# The install files provided to us from tt-mlir are all in a single folder. Organize them so that +# everything is where python expects it to be. +def shuffle_wheel(): + whl = "dist/tt_torch-0.1-cp311-cp311-linux_x86_64.whl" + print(f"Reading in {whl}") + with zipfile.ZipFile(whl, "r") as zip_ref: + zip_ref.extractall("wheel_contents") + + print("Moving tt-metal") + metal_dir = os.path.join( + os.getcwd(), "wheel_contents/tt_torch-0.1.data/data/tt-metal" + ) + shutil.copy2(metal_dir, "wheel_contents/") + shutil.rmtree(metal_dir) + + record_path = "wheel_contents/tt_torch-0.1.dist-info/RECORD" + print("Populating RECORD.") + with open(record_path, "w") as record_file: + for root, _, files in os.walk("wheel_contents"): + for file in files: + file_path = os.path.join(root, file) + rel_path = os.path.relpath(file_path, "wheel_contents") + + # Skip RECORD itself (it should be left without a hash) + if rel_path.endswith("RECORD"): + record_file.write(f"{rel_path},,\n") + continue + + # Compute the hash + hasher = hashlib.sha256() + with open(file_path, "rb") as f: + hasher.update(f.read()) + hash_b64 = hasher.digest().hex() + + # Get file size + file_size = os.path.getsize(file_path) + + # Write to RECORD + record_file.write(f"{rel_path},sha256={hash_b64},{file_size}\n") + + print("Creating wheel.") + with zipfile.ZipFile(whl, "w", zipfile.ZIP_DEFLATED) as zipf: + for root, _, files in os.walk("wheel_contents"): + for file in files: + file_path = os.path.join(root, file) + rel_path = os.path.relpath(file_path, "wheel_contents") + zipf.write(file_path, rel_path) + + shutil.rmtree("wheel_contents", ignore_errors=True) class install_metal_libs(install_lib): @@ -65,3 +119,5 @@ def run(self): "numpy", ], ) + +shuffle_wheel() diff --git a/third_party/CMakeLists.txt b/third_party/CMakeLists.txt index d904c1f0..046df3cb 100644 --- a/third_party/CMakeLists.txt +++ b/third_party/CMakeLists.txt @@ -74,7 +74,8 @@ else() if(DEFINED ENV{WHEEL_BUILD_CUSTOM_INSTALL_PATH}) file(MAKE_DIRECTORY ${CMAKE_INSTALL_PREFIX}/lib/python3.11/site-packages) - install(DIRECTORY ${TTMLIR_INSTALL_PREFIX}/ DESTINATION "${CMAKE_INSTALL_PREFIX}/lib/python3.11/site-packages" USE_SOURCE_PERMISSIONS) + install(DIRECTORY ${TTMLIR_INSTALL_PREFIX}/ DESTINATION "${CMAKE_INSTALL_PREFIX}" USE_SOURCE_PERMISSIONS) + # install(DIRECTORY ${TTMLIR_INSTALL_PREFIX}/ DESTINATION "${CMAKE_INSTALL_PREFIX}/lib/python3.11/site-packages" USE_SOURCE_PERMISSIONS) else() install(DIRECTORY ${TTMLIR_INSTALL_PREFIX}/ DESTINATION "${CMAKE_INSTALL_PREFIX}" USE_SOURCE_PERMISSIONS) endif() diff --git a/tt_torch/__init__.py b/tt_torch/__init__.py index 65b98435..37f51711 100644 --- a/tt_torch/__init__.py +++ b/tt_torch/__init__.py @@ -11,8 +11,14 @@ mp.set_start_method("forkserver", force=True) import os +import sys import importlib.util +# make sure, venv/lib is in the LD_LIBRARY_PATH +lib_path = os.path.join(os.environ["VIRTUAL_ENV"], "lib") +sys.path.append(lib_path) +os.environ["LD_LIBRARY_PATH"] = lib_path + os.pathsep + os.environ["LD_LIBRARY_PATH"] + # find the tt-metal directory, it can either be in the venv if installed from a wheel or in the third_party source tree package_name = "tt-metal" spec = importlib.util.find_spec(package_name) diff --git a/tt_torch/csrc/CMakeLists.txt b/tt_torch/csrc/CMakeLists.txt index e5fee8ec..57844dec 100644 --- a/tt_torch/csrc/CMakeLists.txt +++ b/tt_torch/csrc/CMakeLists.txt @@ -13,12 +13,7 @@ add_dependencies(TT_TORCH_MLIR ) set_target_properties(TT_TORCH_MLIR PROPERTIES COMPILE_FLAGS "-fno-rtti") -if(DEFINED ENV{WHEEL_BUILD_CUSTOM_INSTALL_PATH}) - file(MAKE_DIRECTORY ${CMAKE_INSTALL_PREFIX}/lib/python3.11/site-packages/tt_torch) - install (TARGETS TT_TORCH_MLIR LIBRARY DESTINATION ${CMAKE_INSTALL_PREFIX}/lib/python3.11/site-packages/tt_torch) -else() - install (TARGETS TT_TORCH_MLIR LIBRARY DESTINATION ${CMAKE_INSTALL_PREFIX}/lib) -endif() +install (TARGETS TT_TORCH_MLIR LIBRARY DESTINATION ${CMAKE_INSTALL_PREFIX}/lib) target_include_directories(TT_TORCH_MLIR PUBLIC ${PROJECT_SOURCE_DIR}/third_party/tt-mlir/src/tt-mlir-build/include @@ -86,9 +81,4 @@ set_target_properties(${TARGET_NAME} INSTALL_RPATH "$ORIGIN" ) -if(DEFINED ENV{WHEEL_BUILD_CUSTOM_INSTALL_PATH}) - file(MAKE_DIRECTORY ${CMAKE_INSTALL_PREFIX}/lib/python3.11/site-packages/tt_torch) - install (TARGETS ${TARGET_NAME} LIBRARY DESTINATION ${CMAKE_INSTALL_PREFIX}/lib/python3.11/site-packages/tt_torch) -else() - install (TARGETS ${TARGET_NAME} LIBRARY DESTINATION ${CMAKE_INSTALL_PREFIX}/lib) -endif() +install (TARGETS ${TARGET_NAME} LIBRARY DESTINATION ${CMAKE_INSTALL_PREFIX}/lib) diff --git a/tt_torch/tools/__init__.py b/tt_torch/tools/__init__.py new file mode 100644 index 00000000..e69de29b