Demoing zero-copy save. (#567)

* Demoing zero-copy save.

* Fixing clippy issue.

* Some cleanup.

* Sdist doesn't require feature ?

* Incorrect clean.

* Clippy ?

* Sanity check

* Fixing the doc builder?

* Using pre-commit for quality.

* This should work.

* Clippy variant.

* pyfeature typo.

* Bypassing the necessity for an env ?

* Remove maturin.

* BigEndian fix.

* Only black.

* We need to check both features.

* ??

* Asking for readonly is not possible.

* Before the build.

* Fixing byte-endian?

Author: Narsil

.github/workflows/build_documentation.yml

@@ -11,6 +11,8 @@ on:
 jobs:
   build:
     uses: huggingface/doc-builder/.github/workflows/build_main_documentation.yml@main
+    env:
+      MATURIN_PEP517_ARGS: "--features py311,pyo3/extension-module"
     with:
       commit_sha: ${{ github.sha }}
       package: safetensors

.github/workflows/build_pr_documentation.yml

@@ -14,6 +14,8 @@ concurrency:
 jobs:
   build:
     uses: huggingface/doc-builder/.github/workflows/build_pr_documentation.yml@main
+    env:
+      MATURIN_PEP517_ARGS: "--features py311,pyo3/extension-module"
     with:
       commit_sha: ${{ github.event.pull_request.head.sha }}
       pr_number: ${{ github.event.number }}

.github/workflows/python-release.yml

@@ -23,6 +23,7 @@ jobs:
     runs-on: ${{ matrix.platform.runner }}
     strategy:
       matrix:
+        pyfeature: ["py38", "py311"]
         platform:
           - runner: ubuntu-latest
             target: x86_64
@@ -45,19 +46,20 @@ jobs:
         uses: PyO3/maturin-action@v1
         with:
           target: ${{ matrix.platform.target }}
-          args: --release --out dist --manifest-path bindings/python/Cargo.toml
+          args: --release --out dist --manifest-path bindings/python/Cargo.toml --features pyo3/extension-module,${{ matrix.pyfeature }}
           sccache: 'true'
           manylinux: auto
       - name: Upload wheels
         uses: actions/upload-artifact@v4
         with:
-          name: wheels-linux-${{ matrix.platform.target }}
+          name: wheels-linux-${{ matrix.platform.target }}-${{ matrix.pyfeature }}
           path: dist
 
   musllinux:
     runs-on: ${{ matrix.platform.runner }}
     strategy:
       matrix:
+        pyfeature: ["py38", "py311"]
         platform:
           - runner: ubuntu-latest
             target: x86_64
@@ -76,19 +78,20 @@ jobs:
         uses: PyO3/maturin-action@v1
         with:
           target: ${{ matrix.platform.target }}
-          args: --release --out dist --manifest-path bindings/python/Cargo.toml
+          args: --release --out dist --manifest-path bindings/python/Cargo.toml --features pyo3/extension-module,${{ matrix.pyfeature }}
           sccache: 'true'
           manylinux: musllinux_1_2
       - name: Upload wheels
         uses: actions/upload-artifact@v4
         with:
-          name: wheels-musllinux-${{ matrix.platform.target }}
+          name: wheels-musllinux-${{ matrix.platform.target }}-${{ matrix.pyfeature }}
           path: dist
 
   windows:
     runs-on: ${{ matrix.platform.runner }}
     strategy:
       matrix:
+        pyfeature: ["py38", "py311"]
         platform:
           - runner: windows-latest
             target: x64
@@ -104,18 +107,19 @@ jobs:
         uses: PyO3/maturin-action@v1
         with:
           target: ${{ matrix.platform.target }}
-          args: --release --out dist --manifest-path bindings/python/Cargo.toml
+          args: --release --out dist --manifest-path bindings/python/Cargo.toml --features pyo3/extension-module,${{ matrix.pyfeature }}
           sccache: 'true'
       - name: Upload wheels
         uses: actions/upload-artifact@v4
         with:
-          name: wheels-windows-${{ matrix.platform.target }}
+          name: wheels-windows-${{ matrix.platform.target }}-${{ matrix.pyfeature }}
           path: dist
 
   macos:
     runs-on: ${{ matrix.platform.runner }}
     strategy:
       matrix:
+        pyfeature: ["py38", "py311"]
         platform:
           - runner: macos-13
             target: x86_64
@@ -130,12 +134,12 @@ jobs:
         uses: PyO3/maturin-action@v1
         with:
           target: ${{ matrix.platform.target }}
-          args: --release --out dist --manifest-path bindings/python/Cargo.toml
+          args: --release --out dist --manifest-path bindings/python/Cargo.toml --features pyo3/extension-module,${{ matrix.pyfeature }}
           sccache: 'true'
       - name: Upload wheels
         uses: actions/upload-artifact@v4
         with:
-          name: wheels-macos-${{ matrix.platform.target }}
+          name: wheels-macos-${{ matrix.platform.target }}-${{ matrix.pyfeature }}
           path: dist
 
   sdist:

.github/workflows/python.yml

@@ -7,12 +7,14 @@ jobs:
   build_and_test:
     name: Check everything builds & tests
     runs-on: ${{ matrix.os }}
+    env: 
+      MATURIN_PEP517_ARGS: "--features ${{ matrix.version.pyfeature }},pyo3/extension-module"
     strategy:
       matrix:
         os: [ubuntu-latest, macos-13, windows-latest]
         # Lowest and highest, no version specified so that 
         # new releases get automatically tested against
-        version: [{torch: torch==1.10, python: "3.8"}, {torch: torch, python: "3.12"}]
+        version: [{torch: torch==1.10, python: "3.8", pyfeature: "py38"}, {torch: torch, python: "3.12", pyfeature: "py311"}]
           # TODO this would include macos ARM target.
           # however jax has an illegal instruction issue
           # that exists only in CI (probably difference in instruction support).
@@ -52,14 +54,14 @@ jobs:
         run: cargo fmt -- --check
 
       - name: Lint with Clippy
-        run: cargo clippy --all-targets --all-features -- -D warnings
+        run: | 
+          cargo clippy --features ${{ matrix.version.pyfeature }} -- -D warnings
 
       - name: Run Audit
         run: cargo audit -D warnings
 
       - name: Install
         run: |
-          pip install -U pip
           pip install .[numpy,tensorflow]
           pip install ${{ matrix.version.torch }}
 
@@ -82,7 +84,7 @@ jobs:
 
       - name: Run tests
         run: |
-          cargo test
+          cargo test --features ${{ matrix.version.pyfeature }}
           pip install .[testing]
           pytest -sv tests/
 

.pre-commit-config.yaml

@@ -27,6 +27,8 @@ repos:
           [
             "--manifest-path",
             "bindings/python/Cargo.toml",
+            "--features",
+            "py311",
             "--all-targets",
             "--",
             "-Dwarnings",

Dockerfile.s390x.test

@@ -11,6 +11,7 @@ RUN /root/miniconda3/bin/pip install -U pip pytest
 COPY . .
 SHELL ["/bin/bash", "-c"]
 WORKDIR /safetensors/bindings/python/
+ENV MATURIN_PEP517_ARGS="--features py311,pyo3/extension-module"
 RUN source /root/.cargo/env && /root/miniconda3/bin/pip install -e .
 RUN /root/miniconda3/bin/pytest -sv tests/test_pt_* tests/test_simple.py
 # RUN /root/miniconda3/bin/python -c 'from huggingface_hub import hf_hub_download; filename = hf_hub_download("roberta-base", "model.safetensors"); from safetensors.torch import load_file; weights = load_file(filename); assert weights["roberta.embeddings.position_embeddings.weight"][0][0].abs().item() > 1e-10'

bindings/python/Cargo.toml

@@ -9,8 +9,12 @@ rust-version = "1.74"
 name = "safetensors_rust"
 crate-type = ["cdylib"]
 
+[features]
+py38 = ["pyo3/abi3-py38"]
+py311 = ["pyo3/abi3-py311"]
+
 [dependencies]
-pyo3 = { version = "0.23", features = ["abi3", "abi3-py38"] }
+pyo3 = { version = "0.23", features = ["abi3"] }
 memmap2 = "0.9"
 serde_json = "1.0"
 

bindings/python/py_src/safetensors/torch.py

@@ -128,7 +128,10 @@ def _remove_duplicate_names(
 
 
 def save_model(
-    model: torch.nn.Module, filename: str, metadata: Optional[Dict[str, str]] = None, force_contiguous: bool = True
+    model: torch.nn.Module,
+    filename: str,
+    metadata: Optional[Dict[str, str]] = None,
+    force_contiguous: bool = True,
 ):
     """
     Saves a given torch model to specified filename.
@@ -174,7 +177,10 @@ def save_model(
 
 
 def load_model(
-    model: torch.nn.Module, filename: Union[str, os.PathLike], strict: bool = True, device: Union[str, int] = "cpu"
+    model: torch.nn.Module,
+    filename: Union[str, os.PathLike],
+    strict: bool = True,
+    device: Union[str, int] = "cpu",
 ) -> Tuple[List[str], List[str]]:
     """
     Loads a given filename onto a torch model.
@@ -402,7 +408,7 @@ def _view2torch(safeview) -> Dict[str, torch.Tensor]:
     return result
 
 
-def _tobytes(tensor: torch.Tensor, name: str) -> bytes:
+def _tobytes(tensor: torch.Tensor, name: str) -> Union[memoryview, bytes]:
     if tensor.layout != torch.strided:
         raise ValueError(
             f"You are trying to save a sparse tensor: `{name}` which this library does not support."
@@ -456,8 +462,11 @@ def _tobytes(tensor: torch.Tensor, name: str) -> bytes:
         }
         npdtype = NPDTYPES[tensor.dtype]
         # Not in place as that would potentially modify a live running model
-        data = data.view(npdtype).byteswap(inplace=False)
-    return data.tobytes()
+        data = data.view(npdtype).byteswap(inplace=False).view(np.uint8)
+    if sys.version_info >= (3, 11):
+        return data.data
+    else:
+        return data.tobytes()
 
 
 def _flatten(tensors: Dict[str, torch.Tensor]) -> Dict[str, Dict[str, Any]]:

bindings/python/src/lib.rs

@@ -1,5 +1,7 @@
 #![deny(missing_docs)]
 //! Dummy doc
+#[cfg(any(feature = "py38", feature = "py311"))]
+mod view;
 use memmap2::{Mmap, MmapOptions};
 use pyo3::exceptions::{PyException, PyFileNotFoundError};
 use pyo3::prelude::*;
@@ -10,94 +12,27 @@ use pyo3::Bound as PyBound;
 use pyo3::{intern, PyErr};
 use safetensors::slice::TensorIndexer;
 use safetensors::tensor::{Dtype, Metadata, SafeTensors, TensorInfo, TensorView};
-use safetensors::View;
-use std::borrow::Cow;
 use std::collections::HashMap;
 use std::fs::File;
 use std::iter::FromIterator;
 use std::ops::Bound;
 use std::path::PathBuf;
 use std::sync::Arc;
+#[cfg(any(feature = "py38", feature = "py311"))]
+use view::prepare;
 
 static TORCH_MODULE: GILOnceCell<Py<PyModule>> = GILOnceCell::new();
 static NUMPY_MODULE: GILOnceCell<Py<PyModule>> = GILOnceCell::new();
 static TENSORFLOW_MODULE: GILOnceCell<Py<PyModule>> = GILOnceCell::new();
 static FLAX_MODULE: GILOnceCell<Py<PyModule>> = GILOnceCell::new();
 static MLX_MODULE: GILOnceCell<Py<PyModule>> = GILOnceCell::new();
 
-struct PyView<'a> {
-    shape: Vec<usize>,
-    dtype: Dtype,
-    data: PyBound<'a, PyBytes>,
-    data_len: usize,
-}
-
-impl View for &PyView<'_> {
-    fn data(&self) -> std::borrow::Cow<[u8]> {
-        Cow::Borrowed(self.data.as_bytes())
-    }
-    fn shape(&self) -> &[usize] {
-        &self.shape
-    }
-    fn dtype(&self) -> Dtype {
-        self.dtype
-    }
-    fn data_len(&self) -> usize {
-        self.data_len
-    }
-}
-
-fn prepare(tensor_dict: HashMap<String, PyBound<PyDict>>) -> PyResult<HashMap<String, PyView>> {
-    let mut tensors = HashMap::with_capacity(tensor_dict.len());
-    for (tensor_name, tensor_desc) in &tensor_dict {
-        let shape: Vec<usize> = tensor_desc
-            .get_item("shape")?
-            .ok_or_else(|| SafetensorError::new_err(format!("Missing `shape` in {tensor_desc:?}")))?
-            .extract()?;
-        let pydata: PyBound<PyAny> = tensor_desc.get_item("data")?.ok_or_else(|| {
-            SafetensorError::new_err(format!("Missing `data` in {tensor_desc:?}"))
-        })?;
-        // Make sure it's extractable first.
-        let data: &[u8] = pydata.extract()?;
-        let data_len = data.len();
-        let data: PyBound<PyBytes> = pydata.extract()?;
-        let pydtype = tensor_desc.get_item("dtype")?.ok_or_else(|| {
-            SafetensorError::new_err(format!("Missing `dtype` in {tensor_desc:?}"))
-        })?;
-        let dtype: String = pydtype.extract()?;
-        let dtype = match dtype.as_ref() {
-            "bool" => Dtype::BOOL,
-            "int8" => Dtype::I8,
-            "uint8" => Dtype::U8,
-            "int16" => Dtype::I16,
-            "uint16" => Dtype::U16,
-            "int32" => Dtype::I32,
-            "uint32" => Dtype::U32,
-            "int64" => Dtype::I64,
-            "uint64" => Dtype::U64,
-            "float16" => Dtype::F16,
-            "float32" => Dtype::F32,
-            "float64" => Dtype::F64,
-            "bfloat16" => Dtype::BF16,
-            "float8_e4m3fn" => Dtype::F8_E4M3,
-            "float8_e5m2" => Dtype::F8_E5M2,
-            dtype_str => {
-                return Err(SafetensorError::new_err(format!(
-                    "dtype {dtype_str} is not covered",
-                )));
-            }
-        };
-
-        let tensor = PyView {
-            shape,
-            dtype,
-            data,
-            data_len,
-        };
-        tensors.insert(tensor_name.to_string(), tensor);
-    }
-    Ok(tensors)
-}
+#[cfg(not(any(feature = "py38", feature = "py311")))]
+compile_error!(
+    "At least one python version must be enabled, use `maturin develop --features py311,pyo3/extension-module`"
+);
+#[cfg(all(feature = "py38", feature = "py311"))]
+compile_error!("Only one python version must be enabled");
 
 /// Serializes raw data.
 ///