diff --git a/build2cmake/src/main.rs b/build2cmake/src/main.rs index 54c0619c..f8e0d0b5 100644 --- a/build2cmake/src/main.rs +++ b/build2cmake/src/main.rs @@ -20,6 +20,8 @@ use config::{Backend, Build, BuildCompat}; mod fileset; use fileset::FileSet; +mod metadata; + mod version; #[derive(Parser, Debug)] diff --git a/build2cmake/src/metadata.rs b/build2cmake/src/metadata.rs new file mode 100644 index 00000000..a8a48601 --- /dev/null +++ b/build2cmake/src/metadata.rs @@ -0,0 +1,15 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Deserialize, Serialize)] +#[serde(deny_unknown_fields, rename_all = "kebab-case")] +pub struct Metadata { + python_depends: Vec, +} + +impl Metadata { + pub fn new(python_depends: impl Into>) -> Self { + Self { + python_depends: python_depends.into(), + } + } +} diff --git a/build2cmake/src/templates/windows.cmake b/build2cmake/src/templates/windows.cmake index 6685f97f..a052af3d 100644 --- a/build2cmake/src/templates/windows.cmake +++ b/build2cmake/src/templates/windows.cmake @@ -161,6 +161,11 @@ function(add_local_install_target TARGET_NAME PACKAGE_NAME BUILD_VARIANT_NAME) ${PYTHON_FILES} ${LOCAL_INSTALL_DIR}/ + # Copy metadata.json if it exists + COMMAND ${CMAKE_COMMAND} -E copy_if_different + ${CMAKE_SOURCE_DIR}/metadata.json + ${LOCAL_INSTALL_DIR}/ + COMMENT "Copying shared library and Python files to ${LOCAL_INSTALL_DIR}" COMMAND_EXPAND_LISTS ) diff --git a/build2cmake/src/torch/common.rs b/build2cmake/src/torch/common.rs index e3eea573..f98a97ba 100644 --- a/build2cmake/src/torch/common.rs +++ b/build2cmake/src/torch/common.rs @@ -3,6 +3,7 @@ use itertools::Itertools; use minijinja::{context, Environment}; use crate::config::{Backend, General}; +use crate::metadata::Metadata; use crate::FileSet; pub fn write_pyproject_toml( @@ -32,3 +33,18 @@ pub fn write_pyproject_toml( Ok(()) } + +pub fn write_metadata(backend: Backend, general: &General, file_set: &mut FileSet) -> Result<()> { + let writer = file_set.entry("metadata.json"); + + let python_depends = general + .python_depends() + .chain(general.backend_python_depends(backend)) + .collect::>>()?; + + let metadata = Metadata::new(python_depends); + + serde_json::to_writer(writer, &metadata)?; + + Ok(()) +} diff --git a/build2cmake/src/torch/cpu.rs b/build2cmake/src/torch/cpu.rs index 56f5f0b2..b0864b7b 100644 --- a/build2cmake/src/torch/cpu.rs +++ b/build2cmake/src/torch/cpu.rs @@ -8,6 +8,7 @@ use super::{common::write_pyproject_toml, kernel_ops_identifier}; use crate::{ config::{Backend, Build, Kernel, Torch}, fileset::FileSet, + torch::common::write_metadata, version::Version, }; @@ -52,6 +53,8 @@ pub fn write_torch_ext_cpu( write_torch_registration_macros(&mut file_set)?; + write_metadata(Backend::Cpu, &build.general, &mut file_set)?; + Ok(file_set) } diff --git a/build2cmake/src/torch/cuda.rs b/build2cmake/src/torch/cuda.rs index d9753a79..b80ca433 100644 --- a/build2cmake/src/torch/cuda.rs +++ b/build2cmake/src/torch/cuda.rs @@ -10,6 +10,7 @@ use minijinja::{context, Environment}; use super::common::write_pyproject_toml; use super::kernel_ops_identifier; use crate::config::{Backend, Build, Dependency, Kernel, Torch}; +use crate::torch::common::write_metadata; use crate::version::Version; use crate::FileSet; @@ -65,6 +66,8 @@ pub fn write_torch_ext_cuda( write_torch_registration_macros(&mut file_set)?; + write_metadata(backend, &build.general, &mut file_set)?; + Ok(file_set) } diff --git a/build2cmake/src/torch/metal.rs b/build2cmake/src/torch/metal.rs index 1af2228b..aa961400 100644 --- a/build2cmake/src/torch/metal.rs +++ b/build2cmake/src/torch/metal.rs @@ -8,6 +8,7 @@ use super::{common::write_pyproject_toml, kernel_ops_identifier}; use crate::{ config::{Backend, Build, Kernel, Torch}, fileset::FileSet, + torch::common::write_metadata, version::Version, }; @@ -54,6 +55,8 @@ pub fn write_torch_ext_metal( write_torch_registration_macros(&mut file_set)?; + write_metadata(Backend::Metal, &build.general, &mut file_set)?; + Ok(file_set) } diff --git a/build2cmake/src/torch/noarch.rs b/build2cmake/src/torch/noarch.rs index 902da62b..85caf080 100644 --- a/build2cmake/src/torch/noarch.rs +++ b/build2cmake/src/torch/noarch.rs @@ -7,7 +7,7 @@ use minijinja::{context, Environment}; use crate::{ config::{Backend, Build, General, Torch}, fileset::FileSet, - torch::kernel_ops_identifier, + torch::{common::write_metadata, kernel_ops_identifier}, }; pub fn write_torch_ext_noarch( @@ -30,6 +30,8 @@ pub fn write_torch_ext_noarch( &mut file_set, )?; + write_metadata(backend, &build.general, &mut file_set)?; + Ok(file_set) } diff --git a/build2cmake/src/torch/xpu.rs b/build2cmake/src/torch/xpu.rs index f9da379a..126feb08 100644 --- a/build2cmake/src/torch/xpu.rs +++ b/build2cmake/src/torch/xpu.rs @@ -9,6 +9,7 @@ use minijinja::{context, Environment}; use super::common::write_pyproject_toml; use super::kernel_ops_identifier; use crate::config::{Backend, Build, Dependency, Kernel, Torch}; +use crate::torch::common::write_metadata; use crate::version::Version; use crate::FileSet; @@ -53,6 +54,8 @@ pub fn write_torch_ext_xpu( write_torch_registration_macros(&mut file_set)?; + write_metadata(Backend::Xpu, &build.general, &mut file_set)?; + Ok(file_set) } diff --git a/lib/torch-extension/arch.nix b/lib/torch-extension/arch.nix index 3d4c58af..fc6a8bf1 100644 --- a/lib/torch-extension/arch.nix +++ b/lib/torch-extension/arch.nix @@ -83,12 +83,6 @@ let moduleName = builtins.replaceStrings [ "-" ] [ "_" ] kernelName; - metadata = builtins.toJSON { - python-depends = pythonDeps; - }; - - metadataFile = writeText "metadata.json" metadata; - # On Darwin, we need the host's xcrun for `xcrun metal` to compile Metal shaders. # It's not supported by the nixpkgs shim. xcrunHost = writeScriptBin "xcrunHost" '' @@ -255,7 +249,7 @@ stdenv.mkDerivation (prevAttrs: { mkdir $out/${moduleName} cp ${./compat.py} $out/${moduleName}/__init__.py - cp ${metadataFile} $out/metadata.json + cp ../metadata.json $out/ '' + (lib.optionalString (stripRPath && stdenv.hostPlatform.isLinux)) '' find $out/ -name '*.so' \ diff --git a/lib/torch-extension/no-arch.nix b/lib/torch-extension/no-arch.nix index 36c5d87c..83ebf7b6 100644 --- a/lib/torch-extension/no-arch.nix +++ b/lib/torch-extension/no-arch.nix @@ -51,10 +51,6 @@ let ++ resolveBackendPythonDeps buildConfig.backend backendPythonDeps ++ [ torch ]; moduleName = builtins.replaceStrings [ "-" ] [ "_" ] kernelName; - metadata = builtins.toJSON { - python-depends = pythonDeps; - }; - metadataFile = writeText "metadata.json" metadata; metalSupport = buildConfig.metal or false; in @@ -94,7 +90,7 @@ stdenv.mkDerivation (prevAttrs: { cp -r torch-ext/${moduleName}/* $out/ mkdir $out/${moduleName} cp ${./compat.py} $out/${moduleName}/__init__.py - cp ${metadataFile} $out/metadata.json + cp metadata.json $out/ ''; doInstallCheck = true;