python312Packages.tensorflow: rename generated wheel from tensorflow_cpu to tensorflow
This commit is contained in:
Gaetan Lepage
@@ -54,7 +54,7 @@ let
|
||||
isCudaJetson = cudaSupport && cudaPackages.cudaFlags.isJetsonBuild;
|
||||
isCudaX64 = cudaSupport && stdenv.hostPlatform.isx86_64;
|
||||
in
|
||||
buildPythonPackage {
|
||||
buildPythonPackage rec {
|
||||
pname = "tensorflow" + lib.optionalString cudaSupport "-gpu";
|
||||
version = packages."${"version" + lib.optionalString isCudaJetson "_jetson"}";
|
||||
format = "wheel";
|
||||
@@ -141,73 +141,97 @@ buildPythonPackage {
|
||||
popd
|
||||
'';
|
||||
|
||||
# Note that we need to run *after* the fixup phase because the
|
||||
# libraries are loaded at runtime. If we run in preFixup then
|
||||
# patchelf --shrink-rpath will remove the cuda libraries.
|
||||
postFixup =
|
||||
let
|
||||
# rpaths we only need to add if CUDA is enabled.
|
||||
cudapaths = lib.optionals cudaSupport [
|
||||
cudatoolkit.out
|
||||
cudatoolkit.lib
|
||||
cudnn
|
||||
];
|
||||
# When using the cpu-only wheel, the final package will be named `tensorflow_cpu`.
|
||||
# Then, in each package requiring `tensorflow`, our pythonRuntimeDepsCheck will fail with:
|
||||
# importlib.metadata.PackageNotFoundError: No package metadata was found for tensorflow
|
||||
# Hence, we manually rename the package to `tensorflow`.
|
||||
lib.optionalString ((builtins.match ".*tensorflow_cpu.*" src.url) != null) ''
|
||||
(
|
||||
cd $out/${python.sitePackages}
|
||||
|
||||
libpaths = [
|
||||
(lib.getLib stdenv.cc.cc)
|
||||
zlib
|
||||
];
|
||||
dest="tensorflow-${version}.dist-info"
|
||||
|
||||
rpath = lib.makeLibraryPath (libpaths ++ cudapaths);
|
||||
in
|
||||
lib.optionalString stdenv.hostPlatform.isLinux ''
|
||||
# This is an array containing all the directories in the tensorflow2
|
||||
# package that contain .so files.
|
||||
#
|
||||
# TODO: Create this list programmatically, and remove paths that aren't
|
||||
# actually needed.
|
||||
rrPathArr=(
|
||||
"$out/${python.sitePackages}/tensorflow/"
|
||||
"$out/${python.sitePackages}/tensorflow/core/kernels"
|
||||
"$out/${python.sitePackages}/tensorflow/compiler/mlir/stablehlo/"
|
||||
"$out/${python.sitePackages}/tensorflow/compiler/tf2tensorrt/"
|
||||
"$out/${python.sitePackages}/tensorflow/compiler/tf2xla/ops/"
|
||||
"$out/${python.sitePackages}/tensorflow/include/external/ml_dtypes/"
|
||||
"$out/${python.sitePackages}/tensorflow/lite/experimental/microfrontend/python/ops/"
|
||||
"$out/${python.sitePackages}/tensorflow/lite/python/analyzer_wrapper/"
|
||||
"$out/${python.sitePackages}/tensorflow/lite/python/interpreter_wrapper/"
|
||||
"$out/${python.sitePackages}/tensorflow/lite/python/metrics/"
|
||||
"$out/${python.sitePackages}/tensorflow/lite/python/optimize/"
|
||||
"$out/${python.sitePackages}/tensorflow/python/"
|
||||
"$out/${python.sitePackages}/tensorflow/python/autograph/impl/testing"
|
||||
"$out/${python.sitePackages}/tensorflow/python/client"
|
||||
"$out/${python.sitePackages}/tensorflow/python/data/experimental/service"
|
||||
"$out/${python.sitePackages}/tensorflow/python/framework"
|
||||
"$out/${python.sitePackages}/tensorflow/python/grappler"
|
||||
"$out/${python.sitePackages}/tensorflow/python/lib/core"
|
||||
"$out/${python.sitePackages}/tensorflow/python/lib/io"
|
||||
"$out/${python.sitePackages}/tensorflow/python/platform"
|
||||
"$out/${python.sitePackages}/tensorflow/python/profiler/internal"
|
||||
"$out/${python.sitePackages}/tensorflow/python/saved_model"
|
||||
"$out/${python.sitePackages}/tensorflow/python/util"
|
||||
"$out/${python.sitePackages}/tensorflow/tsl/python/lib/core"
|
||||
"$out/${python.sitePackages}/tensorflow.libs/"
|
||||
"${rpath}"
|
||||
mv tensorflow_cpu-${version}.dist-info "$dest"
|
||||
|
||||
(
|
||||
cd "$dest"
|
||||
|
||||
substituteInPlace METADATA \
|
||||
--replace-fail "tensorflow_cpu" "tensorflow"
|
||||
substituteInPlace RECORD \
|
||||
--replace-fail "tensorflow_cpu" "tensorflow"
|
||||
)
|
||||
)
|
||||
''
|
||||
# Note that we need to run *after* the fixup phase because the
|
||||
# libraries are loaded at runtime. If we run in preFixup then
|
||||
# patchelf --shrink-rpath will remove the cuda libraries.
|
||||
+ (
|
||||
let
|
||||
# rpaths we only need to add if CUDA is enabled.
|
||||
cudapaths = lib.optionals cudaSupport [
|
||||
cudatoolkit.out
|
||||
cudatoolkit.lib
|
||||
cudnn
|
||||
];
|
||||
|
||||
# The the bash array into a colon-separated list of RPATHs.
|
||||
rrPath=$(IFS=$':'; echo "''${rrPathArr[*]}")
|
||||
echo "about to run patchelf with the following rpath: $rrPath"
|
||||
libpaths = [
|
||||
(lib.getLib stdenv.cc.cc)
|
||||
zlib
|
||||
];
|
||||
|
||||
find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
|
||||
echo "about to patchelf $lib..."
|
||||
chmod a+rx "$lib"
|
||||
patchelf --set-rpath "$rrPath" "$lib"
|
||||
${lib.optionalString cudaSupport ''
|
||||
addDriverRunpath "$lib"
|
||||
''}
|
||||
done
|
||||
'';
|
||||
rpath = lib.makeLibraryPath (libpaths ++ cudapaths);
|
||||
in
|
||||
lib.optionalString stdenv.hostPlatform.isLinux ''
|
||||
# This is an array containing all the directories in the tensorflow2
|
||||
# package that contain .so files.
|
||||
#
|
||||
# TODO: Create this list programmatically, and remove paths that aren't
|
||||
# actually needed.
|
||||
rrPathArr=(
|
||||
"$out/${python.sitePackages}/tensorflow/"
|
||||
"$out/${python.sitePackages}/tensorflow/core/kernels"
|
||||
"$out/${python.sitePackages}/tensorflow/compiler/mlir/stablehlo/"
|
||||
"$out/${python.sitePackages}/tensorflow/compiler/tf2tensorrt/"
|
||||
"$out/${python.sitePackages}/tensorflow/compiler/tf2xla/ops/"
|
||||
"$out/${python.sitePackages}/tensorflow/include/external/ml_dtypes/"
|
||||
"$out/${python.sitePackages}/tensorflow/lite/experimental/microfrontend/python/ops/"
|
||||
"$out/${python.sitePackages}/tensorflow/lite/python/analyzer_wrapper/"
|
||||
"$out/${python.sitePackages}/tensorflow/lite/python/interpreter_wrapper/"
|
||||
"$out/${python.sitePackages}/tensorflow/lite/python/metrics/"
|
||||
"$out/${python.sitePackages}/tensorflow/lite/python/optimize/"
|
||||
"$out/${python.sitePackages}/tensorflow/python/"
|
||||
"$out/${python.sitePackages}/tensorflow/python/autograph/impl/testing"
|
||||
"$out/${python.sitePackages}/tensorflow/python/client"
|
||||
"$out/${python.sitePackages}/tensorflow/python/data/experimental/service"
|
||||
"$out/${python.sitePackages}/tensorflow/python/framework"
|
||||
"$out/${python.sitePackages}/tensorflow/python/grappler"
|
||||
"$out/${python.sitePackages}/tensorflow/python/lib/core"
|
||||
"$out/${python.sitePackages}/tensorflow/python/lib/io"
|
||||
"$out/${python.sitePackages}/tensorflow/python/platform"
|
||||
"$out/${python.sitePackages}/tensorflow/python/profiler/internal"
|
||||
"$out/${python.sitePackages}/tensorflow/python/saved_model"
|
||||
"$out/${python.sitePackages}/tensorflow/python/util"
|
||||
"$out/${python.sitePackages}/tensorflow/tsl/python/lib/core"
|
||||
"$out/${python.sitePackages}/tensorflow.libs/"
|
||||
"${rpath}"
|
||||
)
|
||||
|
||||
# The the bash array into a colon-separated list of RPATHs.
|
||||
rrPath=$(IFS=$':'; echo "''${rrPathArr[*]}")
|
||||
echo "about to run patchelf with the following rpath: $rrPath"
|
||||
|
||||
find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
|
||||
echo "about to patchelf $lib..."
|
||||
chmod a+rx "$lib"
|
||||
patchelf --set-rpath "$rrPath" "$lib"
|
||||
${lib.optionalString cudaSupport ''
|
||||
addDriverRunpath "$lib"
|
||||
''}
|
||||
done
|
||||
''
|
||||
);
|
||||
|
||||
# Upstream has a pip hack that results in bin/tensorboard being in both tensorflow
|
||||
# and the propagated input tensorboard, which causes environment collisions.
|
||||
|
||||
Reference in New Issue
Block a user