Skip to content

Commit e32703f

Browse files
takeshi-yoshimuratakeshi-yoshimura
authored
refactor and add publish (#20)
* refactor entire code to add frameworks * add publish * fix bugs in get_multi_cols * mark 0.1.14 Signed-off-by: Takeshi Yoshimura <tyos@jp.ibm.com>
1 parent dd4f237 commit e32703f

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

46 files changed

+3058
-1541
lines changed

‎.github/workflows/publish.yaml‎

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
name: Build wheel and sdist
2+
3+
on:
4+
workflow_dispatch:
5+
release:
6+
types: [published]
7+
8+
jobs:
9+
build-wheels:
10+
name: Build wheel for ${{ matrix.python-version }}
11+
runs-on: ubuntu-latest
12+
strategy:
13+
matrix:
14+
python-version: ["cp39", "cp310", "cp311", "cp312", "cp313"]
15+
16+
steps:
17+
- name: Checkout code
18+
uses: actions/checkout@v4
19+
20+
- name: Set up Python
21+
uses: actions/setup-python@v5
22+
with:
23+
python-version: "3.11"
24+
25+
- name: Install cibuildwheel
26+
run: python -m pip install cibuildwheel setuptools pybind11
27+
- name: Build wheels with cibuildwheel
28+
run: |
29+
cibuildwheel --output-dir wheelhouse
30+
env:
31+
CIBW_BUILD: "${{ matrix.python-version }}-manylinux_x86_64"
32+
CIBW_SKIP: "*-musllinux_* *-win32 *-manylinux_i686"
33+
CIBW_TEST_SKIP: "*"
34+
CIBW_ARCHS: "x86_64"
35+
36+
- name: Upload wheel artifact
37+
uses: actions/upload-artifact@v4
38+
with:
39+
name: wheels-${{ matrix.python-version }}
40+
path: wheelhouse/*.whl
41+
42+
build-sdist:
43+
runs-on: ubuntu-latest
44+
45+
steps:
46+
- uses: actions/checkout@v4
47+
48+
- name: Set up Python
49+
uses: actions/setup-python@v5
50+
with:
51+
python-version: "3.11"
52+
53+
- name: Install build
54+
run: pip install build
55+
56+
- name: Build sdist
57+
run: python -m build --sdist --outdir dist/
58+
59+
- name: Upload sdist artifact
60+
uses: actions/upload-artifact@v4
61+
with:
62+
name: sdist
63+
path: dist/*.tar.gz
64+
65+
upload:
66+
needs: [build-wheels, build-sdist]
67+
runs-on: ubuntu-latest
68+
69+
steps:
70+
- name: Download all wheels
71+
uses: actions/download-artifact@v4
72+
with:
73+
path: dist
74+
- name: Download sdist artifact
75+
uses: actions/download-artifact@v4
76+
with:
77+
name: sdist
78+
path: dist
79+
80+
- name: Flatten all artifacts
81+
run: |
82+
mkdir final_dist
83+
find dist -name '*.whl' -exec cp {} final_dist/ \;
84+
find dist -name '*.tar.gz' -exec cp {} final_dist/ \;
85+
86+
- name: Upload all artifact
87+
uses: actions/upload-artifact@v4
88+
with:
89+
name: final_dist
90+
path: final_dist
91+
92+
# - name: Set up Python
93+
# uses: actions/setup-python@v5
94+
# with:
95+
# python-version: "3.11"
96+
# - name: Publish to PyPI
97+
# run: |
98+
# python -m pip install twine
99+
# twine upload --non-interactive final_dist/*
100+
# env:
101+
# TWINE_USERNAME: __token__
102+
# TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }}

‎.github/workflows/test-paddle.yaml‎

Lines changed: 20 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -32,9 +32,23 @@ jobs:
3232
- name: Install Python dependencies
3333
run: |
3434
python -m pip install --upgrade pip
35-
pip install torch==2.5.1 --index-url https://download.pytorch.org/whl/cpu
35+
tf_ver=4.52
36+
npy_ver=2.2
37+
torch_ver=2.7
38+
if [ "${{ matrix.python-version }}" = "3.9" ]; then
39+
npy_ver=1.26
40+
tf_ver=4.40
41+
torch_ver=2.1
42+
elif [ "${{ matrix.python-version }}" = "3.10" ]; then
43+
torch_ver=2.3
44+
elif [ "${{ matrix.python-version }}" = "3.11" ]; then
45+
torch_ver=2.5
46+
elif [ "${{ matrix.python-version }}" = "3.12" ]; then
47+
torch_ver=2.6
48+
fi
49+
pip install torch==${torch_ver} --index-url https://download.pytorch.org/whl/cpu # transformers requires torch
3650
pip install paddlepaddle==3.0.0
37-
pip install pytest pytest-cov setuptools_scm safetensors transformers==4.52
51+
pip install pytest pytest-cov setuptools_scm safetensors transformers==${tf_ver} numpy==${npy_ver}
3852
- name: Build Package
3953
run: |
4054
pip install .
@@ -43,12 +57,11 @@ jobs:
4357
cd tests
4458
LIBDIR=`python3 -c "import os; os.chdir('/tmp'); import fastsafetensors; print(os.path.dirname(fastsafetensors.__file__))"`
4559
mkdir -p /tmp/pytest-log
60+
export TEST_FASTSAFETENSORS_FRAMEWORK=paddle
4661
COVERAGE_FILE=.coverage_0 pytest -s --cov=${LIBDIR} test_fastsafetensors.py > /tmp/pytest-log/0.log 2>&1
47-
COVERAGE_FILE=.coverage_1 CUDA_VISIBLE_DEVICES="" pytest -s --cov=${LIBDIR} test_fastsafetensors.py > /tmp/pytest-log/1.log 2>&1
48-
COVERAGE_FILE=.coverage_2 torchrun --nnodes=2 --master_addr=0.0.0.0 --master_port=1234 --node_rank=0 --no-python pytest -s --cov=${LIBDIR} test_multi.py > /tmp/pytest-log/2.log 2>&1 &
49-
COVERAGE_FILE=.coverage_3 torchrun --nnodes=2 --master_addr=0.0.0.0 --master_port=1234 --node_rank=1 --no-python pytest -s --cov=${LIBDIR} test_multi.py > /tmp/pytest-log/3.log 2>&1
50-
python -m paddle.distributed.launch --nproc_per_node 2 run_distributed_paddle_test.py -s --cov=${LIBDIR} test_multi_paddle.py
51-
coverage combine .coverage_0 .coverage_1 .coverage_2 .coverage_3 .coverage_4 .coverage_5
62+
COVERAGE_FILE=.coverage_1 WORLD_SIZE=2 python3 -m paddle.distributed.launch --nnodes 2 --master 127.0.0.1:1234 --rank 0 test_multi.py --cov=${LIBDIR} -s test_multi.py > /tmp/pytest-log/1.log 2>&1 & \
63+
COVERAGE_FILE=.coverage_2 WORLD_SIZE=2 python3 -m paddle.distributed.launch --nnodes 2 --master 127.0.0.1:1234 --rank 1 test_multi.py --cov=${LIBDIR} -s test_multi.py > /tmp/pytest-log/2.log 2>&1 && \
64+
coverage combine .coverage_*
5265
coverage html
5366
mv htmlcov /tmp/pytest-log
5467
- name: upload pytest log

‎.github/workflows/test-torch.yaml‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -56,11 +56,14 @@ jobs:
5656
cd tests
5757
LIBDIR=`python3 -c "import os; os.chdir('/tmp'); import fastsafetensors; print(os.path.dirname(fastsafetensors.__file__))"`
5858
mkdir -p /tmp/pytest-log
59+
export TEST_FASTSAFETENSORS_FRAMEWORK=pytorch
5960
COVERAGE_FILE=.coverage_0 pytest -s --cov=${LIBDIR} test_fastsafetensors.py > /tmp/pytest-log/0.log 2>&1
60-
COVERAGE_FILE=.coverage_1 CUDA_VISIBLE_DEVICES="" pytest -s --cov=${LIBDIR} test_fastsafetensors.py > /tmp/pytest-log/1.log 2>&1
61-
COVERAGE_FILE=.coverage_2 torchrun --nnodes=2 --master_addr=0.0.0.0 --master_port=1234 --node_rank=0 --no-python pytest -s --cov=${LIBDIR} test_multi.py > /tmp/pytest-log/2.log 2>&1 &
62-
COVERAGE_FILE=.coverage_3 torchrun --nnodes=2 --master_addr=0.0.0.0 --master_port=1234 --node_rank=1 --no-python pytest -s --cov=${LIBDIR} test_multi.py > /tmp/pytest-log/3.log 2>&1
63-
coverage combine .coverage_0 .coverage_1 .coverage_2 .coverage_3
61+
COVERAGE_FILE=.coverage_1 torchrun --nnodes=1 --master_addr=0.0.0.0 --master_port=1234 --node_rank=0 test_multi.py --cov=${LIBDIR} -s test_multi.py > /tmp/pytest-log/1.log 2>&1
62+
COVERAGE_FILE=.coverage_2 torchrun --nnodes=4 --master_addr=0.0.0.0 --master_port=1234 --node_rank=0 test_multi.py --cov=${LIBDIR} -s test_multi.py > /tmp/pytest-log/2.log 2>&1 &
63+
COVERAGE_FILE=.coverage_3 torchrun --nnodes=4 --master_addr=0.0.0.0 --master_port=1234 --node_rank=1 test_multi.py --cov=${LIBDIR} -s test_multi.py > /tmp/pytest-log/3.log 2>&1 &
64+
COVERAGE_FILE=.coverage_4 torchrun --nnodes=4 --master_addr=0.0.0.0 --master_port=1234 --node_rank=2 test_multi.py --cov=${LIBDIR} -s test_multi.py > /tmp/pytest-log/4.log 2>&1 &
65+
COVERAGE_FILE=.coverage_5 torchrun --nnodes=4 --master_addr=0.0.0.0 --master_port=1234 --node_rank=3 test_multi.py --cov=${LIBDIR} -s test_multi.py > /tmp/pytest-log/5.log 2>&1
66+
coverage combine .coverage_*
6467
coverage html
6568
mv htmlcov /tmp/pytest-log
6669
- name: Upload Pytest log

‎Makefile‎

Lines changed: 45 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -6,24 +6,44 @@ CONCMD := docker
66
ifdef PODMAN
77
CONCMD = podman
88
endif
9-
FST_DIR := $(shell python3 -c "import os; os.chdir('/tmp'); import fastsafetensors; print(os.path.dirname(fastsafetensors.__file__))")
109

1110
.PHONY: install
1211
install:
1312
pip install . --no-cache-dir --no-build-isolation
1413

15-
.PHONY: unittest
14+
.PHONY: unittest unittest-parallel unittest-paddle unittest-paddle-gpu htmlcov
15+
16+
FST_DIR := $(shell python3 -c "import os; os.chdir('/tmp'); import fastsafetensors; print(os.path.dirname(fastsafetensors.__file__))")
17+
1618
unittest:
17-
COVERAGE_FILE=.coverage_0 pytest -s --cov=$(FST_DIR) tests/test_fastsafetensors.py
18-
COVERAGE_FILE=.coverage_1 CUDA_VISIBLE_DEVICES="" pytest -s --cov=$(FST_DIR) tests/test_fastsafetensors.py
19-
COVERAGE_FILE=.coverage_2 torchrun --nnodes=2 --master_addr=0.0.0.0 --master_port=1234 --node_rank=0 --no-python pytest -s --cov=${FST_DIR} tests/test_multi.py > /tmp/2.log 2>&1 &
20-
COVERAGE_FILE=.coverage_3 torchrun --nnodes=2 --master_addr=0.0.0.0 --master_port=1234 --node_rank=1 --no-python pytest -s --cov=${FST_DIR} tests/test_multi.py > /tmp/3.log 2>&1
21-
coverage combine .coverage_0 .coverage_1 .coverage_2 .coverage_3
22-
coverage html
19+
@FST_DIR=$(FST_DIR); \
20+
TEST_FASTSAFETENSORS_FRAMEWORK=torch COVERAGE_FILE=.coverage_0 pytest -s --cov=$(FST_DIR) tests/test_fastsafetensors.py && \
21+
TEST_FASTSAFETENSORS_FRAMEWORK=torch COVERAGE_FILE=.coverage_1 CUDA_VISIBLE_DEVICES="" pytest -s --cov=$(FST_DIR) tests/test_fastsafetensors.py && \
22+
TEST_FASTSAFETENSORS_FRAMEWORK=torch COVERAGE_FILE=.coverage_2 pytest -s --cov=$(FST_DIR) -s tests/test_vllm.py
23+
24+
unittest-parallel:
25+
TEST_FASTSAFETENSORS_FRAMEWORK=torch COVERAGE_FILE=.coverage_3 torchrun --nnodes=4 --master_addr=0.0.0.0 --master_port=1234 --node_rank=0 tests/test_multi.py --cov=$(FST_DIR) -s tests/test_multi.py > /tmp/3.log 2>&1 & \
26+
TEST_FASTSAFETENSORS_FRAMEWORK=torch COVERAGE_FILE=.coverage_4 torchrun --nnodes=4 --master_addr=0.0.0.0 --master_port=1234 --node_rank=1 tests/test_multi.py --cov=$(FST_DIR) -s tests/test_multi.py > /tmp/4.log 2>&1 & \
27+
TEST_FASTSAFETENSORS_FRAMEWORK=torch COVERAGE_FILE=.coverage_5 torchrun --nnodes=4 --master_addr=0.0.0.0 --master_port=1234 --node_rank=2 tests/test_multi.py --cov=$(FST_DIR) -s tests/test_multi.py > /tmp/5.log 2>&1 & \
28+
TEST_FASTSAFETENSORS_FRAMEWORK=torch COVERAGE_FILE=.coverage_6 torchrun --nnodes=4 --master_addr=0.0.0.0 --master_port=1234 --node_rank=3 tests/test_multi.py --cov=$(FST_DIR) -s tests/test_multi.py > /tmp/6.log 2>&1 && \
29+
wait && \
30+
TEST_FASTSAFETENSORS_FRAMEWORK=torch COVERAGE_FILE=.coverage_7 torchrun --nnodes=1 --master_addr=0.0.0.0 --master_port=1234 --node_rank=0 tests/test_multi.py --cov=$(FST_DIR) -s tests/test_multi.py > /tmp/7.log 2>&1 & \
31+
wait
2332

24-
.PHONY: integrationtest
25-
integrationtest:
26-
cd tests && COVERAGE_FILE=.coverage pytest -s test_vllm.py
33+
unittest-paddle:
34+
@FST_DIR=$(FST_DIR); \
35+
TEST_FASTSAFETENSORS_FRAMEWORK=paddle COVERAGE_FILE=.coverage_8 CUDA_VISIBLE_DEVICES="" pytest -s --cov=$(FST_DIR) tests/test_fastsafetensors.py && \
36+
TEST_FASTSAFETENSORS_FRAMEWORK=paddle COVERAGE_FILE=.coverage_9 CUDA_VISIBLE_DEVICES="" WORLD_SIZE=2 python3 -m paddle.distributed.launch --nnodes 2 --master 127.0.0.1:1234 --rank 0 tests/test_multi.py --cov=$(FST_DIR) -s tests/test_multi.py > /tmp/9.log 2>&1 & \
37+
TEST_FASTSAFETENSORS_FRAMEWORK=paddle COVERAGE_FILE=.coverage_10 CUDA_VISIBLE_DEVICES="" WORLD_SIZE=2 python3 -m paddle.distributed.launch --nnodes 2 --master 127.0.0.1:1234 --rank 1 tests/test_multi.py --cov=$(FST_DIR) -s tests/test_multi.py > /tmp/10.log 2>&1 && \
38+
wait
39+
40+
unittest-paddle-gpu:
41+
@FST_DIR=$(FST_DIR); \
42+
TEST_FASTSAFETENSORS_FRAMEWORK=paddle COVERAGE_FILE=.coverage_11 pytest -s --cov=$(FST_DIR) tests/test_fastsafetensors.py
43+
44+
htmlcov:
45+
coverage combine .coverage_* && \
46+
coverage html
2747

2848
.PHONY: builder
2949
builder: Dockerfile.build
@@ -45,7 +65,20 @@ upload:
4565
python3 -m twine upload -u __token__ dist/fastsafetensors-$(shell grep version pyproject.toml | sed -e 's/version = "\([0-9.]\+\)"/\1/g')*
4666

4767
perf/dist:
48-
cd perf && python3 -m build
68+
cd perf && pip install .
69+
70+
.PHONY: format
71+
format:
72+
black8 .
73+
isort .
74+
75+
.PHONY: lint
76+
lint:
77+
black .
78+
isort .
79+
flake8 . --select=E9,F63,F7,F82
80+
mypy . --ignore-missing-imports
4981

82+
.PHONY: clean
5083
clean:
5184
rm -rf dist build fastsafetensors.egg-info

examples/paddle_case/a_paddle.safetensors renamed to examples/a_paddle.safetensors

File renamed without changes.

examples/paddle_case/b_paddle.safetensors renamed to examples/b_paddle.safetensors

File renamed without changes.

‎examples/extract_keys.py‎

Lines changed: 21 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,32 +1,38 @@
1-
import sys
21
import os
3-
import torch
2+
import sys
3+
from typing import Dict, List
4+
5+
from safetensors.torch import load_file
6+
47
from fastsafetensors import SafeTensorsFileLoader, SingleGroup
5-
from safetensors import safe_open
68

79
if __name__ == "__main__":
810
if len(sys.argv) != 2:
911
print("specify a directory containing safetensors files")
1012
sys.exit(1)
11-
loader = SafeTensorsFileLoader(SingleGroup(), torch.device("cpu"), nogds=True)
13+
loader = SafeTensorsFileLoader(SingleGroup(), device="cpu", nogds=True)
1214
input_file_or_dir = sys.argv[1]
13-
src_files = {0: []}
15+
src_files: Dict[int, List[str]] = {0: []}
1416
orig_keys = {}
1517
if os.path.isdir(input_file_or_dir):
1618
for dir, _, files in os.walk(input_file_or_dir):
17-
for filename in files:
18-
if filename.endswith(".safetensors"):
19-
src_files[0].append(f"{dir}/{filename}")
20-
elif os.path.exists(input_file_or_dir) and input_file_or_dir.endswith(".safetensors"):
21-
src_files[0].append(input_file_or_dir)
22-
with safe_open(input_file_or_dir, framework="pytorch") as f:
23-
for key in f.keys():
24-
orig_keys[key] = f.get_tensor(key)
19+
for filename in files:
20+
if filename.endswith(".safetensors"):
21+
src_files[0].append(f"{dir}/{filename}")
22+
elif os.path.exists(input_file_or_dir) and input_file_or_dir.endswith(
23+
".safetensors"
24+
):
25+
src_files[0].append(input_file_or_dir)
26+
orig_keys = load_file(input_file_or_dir)
2527
loader.add_filenames(src_files)
2628
fb = loader.copy_files_to_device()
2729
if len(orig_keys) > 0:
2830
for key in loader.get_keys():
29-
print(f"\"{key}\",{loader.get_shape(key)},{loader.frames[key].data_offsets},{fb.get_tensor(key).dtype},{orig_keys[key].dtype}")
31+
print(
32+
f'"{key}",{loader.get_shape(key)},{loader.frames[key].data_offsets},{fb.get_tensor(key).dtype},{orig_keys[key].dtype}'
33+
)
3034
else:
3135
for key in loader.get_keys():
32-
print(f"\"{key}\",{loader.get_shape(key)},{loader.frames[key].data_offsets},{fb.get_tensor(key).dtype}")
36+
print(
37+
f'"{key}",{loader.get_shape(key)},{loader.frames[key].data_offsets},{fb.get_tensor(key).dtype}'
38+
)

0 commit comments

Comments
 (0)