neuropress 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- neuropress-0.1.0/.editorconfig +30 -0
- neuropress-0.1.0/.github/workflows/ci.yml +67 -0
- neuropress-0.1.0/.github/workflows/publish.yml +183 -0
- neuropress-0.1.0/.gitignore +100 -0
- neuropress-0.1.0/.pre-commit-config.yaml +16 -0
- neuropress-0.1.0/.prospector.yaml +43 -0
- neuropress-0.1.0/.python-version +1 -0
- neuropress-0.1.0/ARCHITECTURE.md +110 -0
- neuropress-0.1.0/Dockerfile.vast +82 -0
- neuropress-0.1.0/LICENSE +201 -0
- neuropress-0.1.0/PKG-INFO +473 -0
- neuropress-0.1.0/README.md +420 -0
- neuropress-0.1.0/benchmarks/bench_tiered_vq.py +501 -0
- neuropress-0.1.0/benchmarks/bench_vptq_distribution.py +207 -0
- neuropress-0.1.0/benchmarks/test_bench_packing.py +187 -0
- neuropress-0.1.0/configs/lowering/auto.yaml +9 -0
- neuropress-0.1.0/configs/lowering/compile.yaml +10 -0
- neuropress-0.1.0/configs/lowering/cuda_jit.yaml +11 -0
- neuropress-0.1.0/configs/lowering/cuda_lut.yaml +10 -0
- neuropress-0.1.0/configs/lowering/cuda_wmma.yaml +10 -0
- neuropress-0.1.0/neuropress/__init__.py +127 -0
- neuropress-0.1.0/neuropress/adapter.py +125 -0
- neuropress-0.1.0/neuropress/calibration/__init__.py +29 -0
- neuropress-0.1.0/neuropress/calibration/data/__init__.py +26 -0
- neuropress-0.1.0/neuropress/calibration/data/audio.py +147 -0
- neuropress-0.1.0/neuropress/calibration/data/llm.py +415 -0
- neuropress-0.1.0/neuropress/calibration/data/registry.py +22 -0
- neuropress-0.1.0/neuropress/calibration/probe/__init__.py +36 -0
- neuropress-0.1.0/neuropress/calibration/probe/channel_max.py +37 -0
- neuropress-0.1.0/neuropress/calibration/probe/collectors.py +71 -0
- neuropress-0.1.0/neuropress/calibration/probe/hessian.py +216 -0
- neuropress-0.1.0/neuropress/calibration/probe/hessian_diag.py +143 -0
- neuropress-0.1.0/neuropress/calibration/probe/min_max.py +27 -0
- neuropress-0.1.0/neuropress/calibration/probe/registry.py +225 -0
- neuropress-0.1.0/neuropress/cloud/__init__.py +1 -0
- neuropress-0.1.0/neuropress/cloud/cli.py +38 -0
- neuropress-0.1.0/neuropress/cloud/direct_runner.py +123 -0
- neuropress-0.1.0/neuropress/cloud/runner.py +1573 -0
- neuropress-0.1.0/neuropress/cloud/tui.py +511 -0
- neuropress-0.1.0/neuropress/cloud/vast_runner.py +2936 -0
- neuropress-0.1.0/neuropress/estimator.py +103 -0
- neuropress-0.1.0/neuropress/estimator_configs.py +629 -0
- neuropress-0.1.0/neuropress/estimators.py +3308 -0
- neuropress-0.1.0/neuropress/eval/__init__.py +26 -0
- neuropress-0.1.0/neuropress/eval/harness.py +346 -0
- neuropress-0.1.0/neuropress/eval/llm_cli.py +1102 -0
- neuropress-0.1.0/neuropress/eval/size_estimate.py +600 -0
- neuropress-0.1.0/neuropress/eval/timing.py +121 -0
- neuropress-0.1.0/neuropress/graph/__init__.py +7 -0
- neuropress-0.1.0/neuropress/graph/builder.py +93 -0
- neuropress-0.1.0/neuropress/graph/graph.py +43 -0
- neuropress-0.1.0/neuropress/graph/node.py +28 -0
- neuropress-0.1.0/neuropress/kernel_lowering.py +357 -0
- neuropress-0.1.0/neuropress/kernels/__init__.py +1006 -0
- neuropress-0.1.0/neuropress/kernels/kmeans_assign.cu +468 -0
- neuropress-0.1.0/neuropress/kernels/q4_linear.cu +274 -0
- neuropress-0.1.0/neuropress/kernels/rms_norm.cu +123 -0
- neuropress-0.1.0/neuropress/kernels/vq_linear.cu +489 -0
- neuropress-0.1.0/neuropress/kernels/vq_linear_lut.cu +211 -0
- neuropress-0.1.0/neuropress/kernels/vq_linear_ntier.cu +699 -0
- neuropress-0.1.0/neuropress/kernels/vq_linear_tiered.cu +632 -0
- neuropress-0.1.0/neuropress/kernels/vq_linear_wmma.cu +557 -0
- neuropress-0.1.0/neuropress/kernels/vq_linear_zero_escape.cu +559 -0
- neuropress-0.1.0/neuropress/meta.py +93 -0
- neuropress-0.1.0/neuropress/methods/__init__.py +67 -0
- neuropress-0.1.0/neuropress/methods/base.py +147 -0
- neuropress-0.1.0/neuropress/methods/correction_mode.py +15 -0
- neuropress-0.1.0/neuropress/methods/gptq.py +558 -0
- neuropress-0.1.0/neuropress/methods/gptq_panel.py +497 -0
- neuropress-0.1.0/neuropress/methods/hqq.py +235 -0
- neuropress-0.1.0/neuropress/methods/smooth_quant.py +250 -0
- neuropress-0.1.0/neuropress/methods/vptq.py +348 -0
- neuropress-0.1.0/neuropress/methods/vptq_analysis.py +926 -0
- neuropress-0.1.0/neuropress/methods/vq_gptq.py +258 -0
- neuropress-0.1.0/neuropress/methods/zp_scale_calib.py +240 -0
- neuropress-0.1.0/neuropress/nn/__init__.py +54 -0
- neuropress-0.1.0/neuropress/nn/backends.py +159 -0
- neuropress-0.1.0/neuropress/nn/np_embedding.py +599 -0
- neuropress-0.1.0/neuropress/nn/np_linear.py +851 -0
- neuropress-0.1.0/neuropress/nn/np_rms_norm.py +180 -0
- neuropress-0.1.0/neuropress/nn/np_vq_embedding.py +289 -0
- neuropress-0.1.0/neuropress/nn/np_vq_linear.py +1423 -0
- neuropress-0.1.0/neuropress/pipeline.py +142 -0
- neuropress-0.1.0/neuropress/plan.py +32 -0
- neuropress-0.1.0/neuropress/preprocessing.py +279 -0
- neuropress-0.1.0/neuropress/qscheme/__init__.py +15 -0
- neuropress-0.1.0/neuropress/qscheme/descriptor.py +77 -0
- neuropress-0.1.0/neuropress/qscheme/group.py +44 -0
- neuropress-0.1.0/neuropress/qscheme/zp_scale.py +144 -0
- neuropress-0.1.0/neuropress/quantization/__init__.py +27 -0
- neuropress-0.1.0/neuropress/quantization/cache.py +211 -0
- neuropress-0.1.0/neuropress/quantization/spans.py +86 -0
- neuropress-0.1.0/neuropress/recipe.py +58 -0
- neuropress-0.1.0/neuropress/recipe_config.py +488 -0
- neuropress-0.1.0/neuropress/t2n/__init__.py +59 -0
- neuropress-0.1.0/neuropress/t2n/_qtensor_patch.py +103 -0
- neuropress-0.1.0/neuropress/t2n/model_io.py +563 -0
- neuropress-0.1.0/neuropress/t2n/packing.py +224 -0
- neuropress-0.1.0/neuropress/t2n/qtensors.py +526 -0
- neuropress-0.1.0/neuropress/t2n/registry.py +952 -0
- neuropress-0.1.0/neuropress/targets.py +28 -0
- neuropress-0.1.0/neuropress/tools/__init__.py +1 -0
- neuropress-0.1.0/neuropress/tools/bench.py +1649 -0
- neuropress-0.1.0/neuropress/tools/vq_probe.py +115 -0
- neuropress-0.1.0/neuropress/transforms/__init__.py +34 -0
- neuropress-0.1.0/neuropress/transforms/cluster_backends.py +1142 -0
- neuropress-0.1.0/neuropress/transforms/rotation.py +452 -0
- neuropress-0.1.0/neuropress/transforms/rotation_surgery.py +154 -0
- neuropress-0.1.0/neuropress/transforms/split_mixed_linear.py +306 -0
- neuropress-0.1.0/neuropress/transforms/vq.py +529 -0
- neuropress-0.1.0/neuropress/transforms/vq_packing.py +1384 -0
- neuropress-0.1.0/neuropress/transforms/weight_only.py +593 -0
- neuropress-0.1.0/pyproject.toml +150 -0
- neuropress-0.1.0/run +72 -0
- neuropress-0.1.0/scripts/bench_assign.py +222 -0
- neuropress-0.1.0/scripts/bench_kmeans.py +317 -0
- neuropress-0.1.0/scripts/bench_mstep.py +169 -0
- neuropress-0.1.0/scripts/benchmark_inference.py +467 -0
- neuropress-0.1.0/scripts/plot_bench.py +324 -0
- neuropress-0.1.0/scripts/profile_layers.py +202 -0
- neuropress-0.1.0/scripts/profile_mstep.py +262 -0
- neuropress-0.1.0/scripts/test_save.py +233 -0
- neuropress-0.1.0/scripts/vast_runner.py +9 -0
- neuropress-0.1.0/tests/__init__.py +0 -0
- neuropress-0.1.0/tests/calibration/__init__.py +0 -0
- neuropress-0.1.0/tests/calibration/test_audio_data.py +334 -0
- neuropress-0.1.0/tests/calibration/test_hessian.py +140 -0
- neuropress-0.1.0/tests/calibration/test_llm_data.py +397 -0
- neuropress-0.1.0/tests/calibration/test_stat_field.py +143 -0
- neuropress-0.1.0/tests/integration/test_gptq_nnef_export.py +147 -0
- neuropress-0.1.0/tests/kernels/__init__.py +0 -0
- neuropress-0.1.0/tests/kernels/test_kmeans_assign_grouped.py +148 -0
- neuropress-0.1.0/tests/methods/__init__.py +0 -0
- neuropress-0.1.0/tests/methods/test_down_proj_smoothing.py +157 -0
- neuropress-0.1.0/tests/methods/test_ecvq_annealing.py +348 -0
- neuropress-0.1.0/tests/methods/test_global_correction_modes.py +337 -0
- neuropress-0.1.0/tests/methods/test_gptq.py +152 -0
- neuropress-0.1.0/tests/methods/test_gptq_panel.py +265 -0
- neuropress-0.1.0/tests/methods/test_gptq_pipeline.py +158 -0
- neuropress-0.1.0/tests/methods/test_hqq.py +498 -0
- neuropress-0.1.0/tests/methods/test_smooth_quant.py +300 -0
- neuropress-0.1.0/tests/methods/test_vptq_analysis.py +192 -0
- neuropress-0.1.0/tests/methods/test_zp_scale.py +212 -0
- neuropress-0.1.0/tests/nn/test_np_embedding.py +149 -0
- neuropress-0.1.0/tests/nn/test_np_rms_norm.py +240 -0
- neuropress-0.1.0/tests/nn/test_np_vq_linear.py +414 -0
- neuropress-0.1.0/tests/t2n/__init__.py +0 -0
- neuropress-0.1.0/tests/t2n/test_model_io.py +352 -0
- neuropress-0.1.0/tests/t2n/test_model_size.py +191 -0
- neuropress-0.1.0/tests/t2n/test_qtensors.py +211 -0
- neuropress-0.1.0/tests/t2n/test_vq_serde.py +171 -0
- neuropress-0.1.0/tests/test_adapter.py +50 -0
- neuropress-0.1.0/tests/test_cache_geometry_validation.py +95 -0
- neuropress-0.1.0/tests/test_calibration.py +99 -0
- neuropress-0.1.0/tests/test_cloud_cli_entrypoint.py +77 -0
- neuropress-0.1.0/tests/test_cloud_runner.py +160 -0
- neuropress-0.1.0/tests/test_cuda_kmeans.py +112 -0
- neuropress-0.1.0/tests/test_eval.py +243 -0
- neuropress-0.1.0/tests/test_graph.py +131 -0
- neuropress-0.1.0/tests/test_per_group_centroids.py +37 -0
- neuropress-0.1.0/tests/test_pipeline.py +95 -0
- neuropress-0.1.0/tests/test_plan.py +80 -0
- neuropress-0.1.0/tests/test_preprocessing.py +289 -0
- neuropress-0.1.0/tests/test_qscheme.py +140 -0
- neuropress-0.1.0/tests/test_readme_references.py +48 -0
- neuropress-0.1.0/tests/test_recipe_config.py +382 -0
- neuropress-0.1.0/tests/test_size_estimate_cli.py +73 -0
- neuropress-0.1.0/tests/test_t2n_register.py +335 -0
- neuropress-0.1.0/tests/test_targets.py +40 -0
- neuropress-0.1.0/tests/test_transforms.py +330 -0
- neuropress-0.1.0/tests/test_vptq_estimator.py +424 -0
- neuropress-0.1.0/tests/test_vptq_freq_sort.py +144 -0
- neuropress-0.1.0/tests/test_weight_only_transforms.py +120 -0
- neuropress-0.1.0/tests/transforms/__init__.py +0 -0
- neuropress-0.1.0/tests/transforms/test_cluster_backend_selection.py +42 -0
- neuropress-0.1.0/tests/transforms/test_cluster_backends.py +333 -0
- neuropress-0.1.0/tests/transforms/test_ntier_packing.py +303 -0
- neuropress-0.1.0/tests/transforms/test_rotation.py +332 -0
- neuropress-0.1.0/tests/transforms/test_rotation_cuda.py +104 -0
- neuropress-0.1.0/tests/transforms/test_rotation_surgery.py +230 -0
- neuropress-0.1.0/tests/transforms/test_tiered_packing.py +296 -0
- neuropress-0.1.0/tests/transforms/test_vq_packing.py +120 -0
- neuropress-0.1.0/tests/transforms/test_vq_packing_extended.py +207 -0
- neuropress-0.1.0/tests/transforms/test_zero_escape_packing.py +370 -0
|
@@ -0,0 +1,30 @@
|
|
|
1
|
+
root = true
|
|
2
|
+
|
|
3
|
+
[*]
|
|
4
|
+
charset = utf-8
|
|
5
|
+
end_of_line = lf
|
|
6
|
+
insert_final_newline = true
|
|
7
|
+
trim_trailing_whitespace = true
|
|
8
|
+
|
|
9
|
+
[*.py]
|
|
10
|
+
indent_style = space
|
|
11
|
+
indent_size = 4
|
|
12
|
+
max_line_length = 88
|
|
13
|
+
|
|
14
|
+
[*.{toml,cfg,ini}]
|
|
15
|
+
indent_style = space
|
|
16
|
+
indent_size = 4
|
|
17
|
+
|
|
18
|
+
[*.{yml,yaml}]
|
|
19
|
+
indent_style = space
|
|
20
|
+
indent_size = 2
|
|
21
|
+
|
|
22
|
+
[*.{json,jsonc}]
|
|
23
|
+
indent_style = space
|
|
24
|
+
indent_size = 2
|
|
25
|
+
|
|
26
|
+
[*.md]
|
|
27
|
+
trim_trailing_whitespace = false
|
|
28
|
+
|
|
29
|
+
[Makefile]
|
|
30
|
+
indent_style = tab
|
|
@@ -0,0 +1,67 @@
|
|
|
1
|
+
name: ci
|
|
2
|
+
|
|
3
|
+
# Lint and test on every push to main and every pull request.
|
|
4
|
+
#
|
|
5
|
+
# The release workflow (publish.yml) verifies that the *artifact* installs and
|
|
6
|
+
# imports; this one verifies that the *code* is correct. Both gate a release:
|
|
7
|
+
# publish.yml runs on a tag, by which point these checks have already run on the
|
|
8
|
+
# commit being tagged.
|
|
9
|
+
|
|
10
|
+
on:
|
|
11
|
+
push:
|
|
12
|
+
branches: [main]
|
|
13
|
+
pull_request:
|
|
14
|
+
|
|
15
|
+
permissions:
|
|
16
|
+
contents: read
|
|
17
|
+
|
|
18
|
+
concurrency:
|
|
19
|
+
group: ci-${{ github.ref }}
|
|
20
|
+
cancel-in-progress: true
|
|
21
|
+
|
|
22
|
+
jobs:
|
|
23
|
+
lint:
|
|
24
|
+
name: ruff
|
|
25
|
+
runs-on: ubuntu-latest
|
|
26
|
+
steps:
|
|
27
|
+
- uses: actions/checkout@v5
|
|
28
|
+
|
|
29
|
+
- uses: astral-sh/setup-uv@v6
|
|
30
|
+
with:
|
|
31
|
+
python-version: "3.13"
|
|
32
|
+
|
|
33
|
+
# Mirrors the blocking half of `./run check`. Pinned so a new ruff release
|
|
34
|
+
# cannot fail CI on a commit that did not change; bump it deliberately.
|
|
35
|
+
- name: Lint
|
|
36
|
+
run: uvx ruff@0.15.2 check .
|
|
37
|
+
|
|
38
|
+
- name: Format
|
|
39
|
+
run: uvx ruff@0.15.2 format --check .
|
|
40
|
+
|
|
41
|
+
test:
|
|
42
|
+
name: pytest
|
|
43
|
+
runs-on: ubuntu-latest
|
|
44
|
+
steps:
|
|
45
|
+
- uses: actions/checkout@v5
|
|
46
|
+
|
|
47
|
+
- uses: astral-sh/setup-uv@v6
|
|
48
|
+
with:
|
|
49
|
+
python-version: "3.13"
|
|
50
|
+
|
|
51
|
+
- name: Install
|
|
52
|
+
run: |
|
|
53
|
+
set -euo pipefail
|
|
54
|
+
uv venv --python 3.13
|
|
55
|
+
# CPU torch first: the default PyPI wheel bundles CUDA (~2.5 GB) and
|
|
56
|
+
# fills the runner disk. Installing it up front leaves the project
|
|
57
|
+
# install below with the requirement already satisfied.
|
|
58
|
+
uv pip install --python .venv torch --index-url https://download.pytorch.org/whl/cpu
|
|
59
|
+
# `cloud` is needed for collection, not just for coverage:
|
|
60
|
+
# tests/test_cloud_runner.py imports neuropress.cloud.vast_runner,
|
|
61
|
+
# which imports rich at module scope.
|
|
62
|
+
uv pip install --python .venv -e ".[dev,cloud]"
|
|
63
|
+
|
|
64
|
+
- name: Test
|
|
65
|
+
# Integration tests need HuggingFace downloads and several minutes; they
|
|
66
|
+
# are additionally gated behind RUN_INTEGRATION inside the module.
|
|
67
|
+
run: .venv/bin/python -m pytest --ignore=tests/integration -q
|
|
@@ -0,0 +1,183 @@
|
|
|
1
|
+
name: publish
|
|
2
|
+
|
|
3
|
+
# Build neuropress and upload it to PyPI.
|
|
4
|
+
#
|
|
5
|
+
# release (published) -> PyPI (the tag must match the package version)
|
|
6
|
+
# workflow_dispatch -> TestPyPI (rehearsal; no tag check)
|
|
7
|
+
#
|
|
8
|
+
# Authentication is PyPI Trusted Publishing (OIDC). No API token is stored in
|
|
9
|
+
# this repository, so there is nothing to rotate or leak. One-time setup, on
|
|
10
|
+
# both pypi.org and test.pypi.org, before the first upload:
|
|
11
|
+
#
|
|
12
|
+
# Your projects -> Publishing -> Add a pending publisher
|
|
13
|
+
# PyPI project name : neuropress
|
|
14
|
+
# Owner : DreamerMind
|
|
15
|
+
# Repository : neuropress
|
|
16
|
+
# Workflow name : publish.yml
|
|
17
|
+
# Environment : pypi (use "testpypi" on test.pypi.org)
|
|
18
|
+
#
|
|
19
|
+
# Cutting a release:
|
|
20
|
+
# 1. bump `version` in pyproject.toml and commit it
|
|
21
|
+
# 2. gh release create v<version> --generate-notes
|
|
22
|
+
#
|
|
23
|
+
# PyPI is append-only: a version number can never be reused, not even after the
|
|
24
|
+
# file is deleted. Rehearse with "Run workflow" against TestPyPI first.
|
|
25
|
+
|
|
26
|
+
on:
|
|
27
|
+
release:
|
|
28
|
+
types: [published]
|
|
29
|
+
workflow_dispatch:
|
|
30
|
+
|
|
31
|
+
permissions:
|
|
32
|
+
contents: read
|
|
33
|
+
|
|
34
|
+
concurrency:
|
|
35
|
+
group: publish-${{ github.ref }}
|
|
36
|
+
cancel-in-progress: false
|
|
37
|
+
|
|
38
|
+
jobs:
|
|
39
|
+
build:
|
|
40
|
+
name: build and verify artifacts
|
|
41
|
+
runs-on: ubuntu-latest
|
|
42
|
+
steps:
|
|
43
|
+
- uses: actions/checkout@v5
|
|
44
|
+
|
|
45
|
+
- uses: astral-sh/setup-uv@v6
|
|
46
|
+
with:
|
|
47
|
+
python-version: "3.13"
|
|
48
|
+
|
|
49
|
+
- name: Build sdist and wheel
|
|
50
|
+
run: uv build
|
|
51
|
+
|
|
52
|
+
- name: Validate package metadata
|
|
53
|
+
run: uvx twine check --strict dist/*
|
|
54
|
+
|
|
55
|
+
- name: Tag must match the built version
|
|
56
|
+
if: github.event_name == 'release'
|
|
57
|
+
run: |
|
|
58
|
+
# Compare as versions, not as strings: a wheel filename is always
|
|
59
|
+
# PEP 440-normalized, so `version = "1.0.0-rc1"` builds as 1.0.0rc1
|
|
60
|
+
# and would never string-match the natural tag v1.0.0-rc1.
|
|
61
|
+
set -euo pipefail
|
|
62
|
+
built=$(basename "$(ls dist/*.whl)" | cut -d- -f2)
|
|
63
|
+
uv run --no-project --with packaging python - "$built" "${GITHUB_REF_NAME#v}" <<'PY'
|
|
64
|
+
import sys
|
|
65
|
+
|
|
66
|
+
from packaging.version import InvalidVersion, Version
|
|
67
|
+
|
|
68
|
+
built, tag = sys.argv[1], sys.argv[2]
|
|
69
|
+
try:
|
|
70
|
+
match = Version(built) == Version(tag)
|
|
71
|
+
except InvalidVersion:
|
|
72
|
+
match = False
|
|
73
|
+
if not match:
|
|
74
|
+
print(
|
|
75
|
+
f"::error::release tag {tag!r} and the built version {built!r} disagree — "
|
|
76
|
+
"correct whichever is wrong (the tag, or `version` in pyproject.toml)"
|
|
77
|
+
)
|
|
78
|
+
raise SystemExit(1)
|
|
79
|
+
print(f"publishing neuropress {built}")
|
|
80
|
+
PY
|
|
81
|
+
|
|
82
|
+
- name: Artifacts must carry no research-only material
|
|
83
|
+
run: |
|
|
84
|
+
# Defence in depth. pyproject.toml already excludes these from both
|
|
85
|
+
# build targets, so a hit means an exclude pattern regressed. Uploads
|
|
86
|
+
# are permanent, so fail closed rather than warn.
|
|
87
|
+
#
|
|
88
|
+
# Build the listing before grepping it: `unzip -l ... | grep` inside an
|
|
89
|
+
# `if` swallows a failing unzip, and a check that cannot read the
|
|
90
|
+
# artifact must not report it clean.
|
|
91
|
+
#
|
|
92
|
+
# `unzip -Z1` lists bare paths. `unzip -l` pads each name into a
|
|
93
|
+
# right-hand column, which leaves a top-level `reports/` preceded by
|
|
94
|
+
# spaces rather than by start-of-line, so the anchor below could never
|
|
95
|
+
# match it and the wheel half of this check would be dead code.
|
|
96
|
+
set -euo pipefail
|
|
97
|
+
wheel=$(ls dist/*.whl)
|
|
98
|
+
sdist=$(ls dist/*.tar.gz)
|
|
99
|
+
listing=$(unzip -Z1 "$wheel"; tar -tzf "$sdist")
|
|
100
|
+
# Scoped to scripts/: a bare `_diag\.py` also matches shipped modules
|
|
101
|
+
# such as neuropress/calibration/probe/hessian_diag.py.
|
|
102
|
+
pattern='research_registry|research_jobs|sync_public|public-manifest'
|
|
103
|
+
pattern="$pattern"'|(^|/)scripts/(qwen35|gemma4)_'
|
|
104
|
+
pattern="$pattern"'|(^|/)(reports|plans|docs|papers|sweeps)/'
|
|
105
|
+
if grep -Ei "$pattern" <<<"$listing"; then
|
|
106
|
+
echo "::error::research-only paths found in the distribution artifacts"
|
|
107
|
+
exit 1
|
|
108
|
+
fi
|
|
109
|
+
echo "artifacts clean: $wheel $sdist"
|
|
110
|
+
|
|
111
|
+
- uses: actions/upload-artifact@v4
|
|
112
|
+
with:
|
|
113
|
+
name: dist
|
|
114
|
+
path: dist/
|
|
115
|
+
|
|
116
|
+
smoke:
|
|
117
|
+
name: built wheel installs and imports
|
|
118
|
+
needs: build
|
|
119
|
+
runs-on: ubuntu-latest
|
|
120
|
+
steps:
|
|
121
|
+
- uses: actions/download-artifact@v4
|
|
122
|
+
with:
|
|
123
|
+
name: dist
|
|
124
|
+
path: dist
|
|
125
|
+
|
|
126
|
+
- uses: astral-sh/setup-uv@v6
|
|
127
|
+
with:
|
|
128
|
+
python-version: "3.13"
|
|
129
|
+
|
|
130
|
+
- name: Install the wheel with base dependencies only
|
|
131
|
+
run: |
|
|
132
|
+
uv venv --python 3.13 .venv
|
|
133
|
+
# torch's default PyPI wheel bundles CUDA (~2.5 GB) and fills the
|
|
134
|
+
# runner disk. The CPU build is all an import check needs.
|
|
135
|
+
uv pip install --python .venv torch --index-url https://download.pytorch.org/whl/cpu
|
|
136
|
+
uv pip install --python .venv dist/*.whl
|
|
137
|
+
|
|
138
|
+
- name: Import from outside the workspace
|
|
139
|
+
run: |
|
|
140
|
+
# Never run this from a directory containing neuropress/: python puts
|
|
141
|
+
# the working directory on sys.path, so the import would resolve to
|
|
142
|
+
# the source tree and prove nothing about the wheel.
|
|
143
|
+
mkdir -p "$RUNNER_TEMP/smoke"
|
|
144
|
+
cd "$RUNNER_TEMP/smoke"
|
|
145
|
+
"$GITHUB_WORKSPACE/.venv/bin/python" - <<'PY'
|
|
146
|
+
from importlib.metadata import version
|
|
147
|
+
|
|
148
|
+
import importlib.util
|
|
149
|
+
|
|
150
|
+
import neuropress
|
|
151
|
+
from neuropress.t2n import NEUROPRESS_COMPRESSION as registry
|
|
152
|
+
|
|
153
|
+
print(f"neuropress {version('neuropress')} from {neuropress.__file__}")
|
|
154
|
+
assert registry, "compression registry is empty"
|
|
155
|
+
|
|
156
|
+
# The invariant is that the research module is absent, not that its
|
|
157
|
+
# entries happen to share a key prefix — a prefix nothing enforces.
|
|
158
|
+
spec = importlib.util.find_spec("neuropress.t2n.research_registry")
|
|
159
|
+
assert spec is None, f"research module reached the wheel: {spec.origin}"
|
|
160
|
+
print(f"{len(registry)} compression entries, research module absent")
|
|
161
|
+
PY
|
|
162
|
+
|
|
163
|
+
publish:
|
|
164
|
+
name: upload to ${{ github.event_name == 'release' && 'PyPI' || 'TestPyPI' }}
|
|
165
|
+
needs: [build, smoke]
|
|
166
|
+
runs-on: ubuntu-latest
|
|
167
|
+
environment:
|
|
168
|
+
name: ${{ github.event_name == 'release' && 'pypi' || 'testpypi' }}
|
|
169
|
+
url: ${{ github.event_name == 'release' && 'https://pypi.org/p/neuropress' || 'https://test.pypi.org/p/neuropress' }}
|
|
170
|
+
permissions:
|
|
171
|
+
id-token: write # mints the OIDC token trusted publishing exchanges for an upload token
|
|
172
|
+
steps:
|
|
173
|
+
- uses: actions/download-artifact@v4
|
|
174
|
+
with:
|
|
175
|
+
name: dist
|
|
176
|
+
path: dist
|
|
177
|
+
|
|
178
|
+
- uses: pypa/gh-action-pypi-publish@release/v1
|
|
179
|
+
with:
|
|
180
|
+
repository-url: ${{ github.event_name == 'release' && 'https://upload.pypi.org/legacy/' || 'https://test.pypi.org/legacy/' }}
|
|
181
|
+
# Rehearsals re-upload the same version repeatedly; a real release
|
|
182
|
+
# must still fail loudly if that version is already on PyPI.
|
|
183
|
+
skip-existing: ${{ github.event_name != 'release' }}
|
|
@@ -0,0 +1,100 @@
|
|
|
1
|
+
# === Python ===
|
|
2
|
+
__pycache__/
|
|
3
|
+
*.py[cod]
|
|
4
|
+
*.pyo
|
|
5
|
+
*.pyd
|
|
6
|
+
*.so
|
|
7
|
+
*.egg
|
|
8
|
+
*.egg-info/
|
|
9
|
+
dist/
|
|
10
|
+
build/
|
|
11
|
+
wheels/
|
|
12
|
+
MANIFEST
|
|
13
|
+
pip-wheel-metadata/
|
|
14
|
+
share/python-wheels/
|
|
15
|
+
|
|
16
|
+
# === Virtual environments ===
|
|
17
|
+
.venv/
|
|
18
|
+
venv/
|
|
19
|
+
env/
|
|
20
|
+
.env/
|
|
21
|
+
|
|
22
|
+
# === uv / pip ===
|
|
23
|
+
uv.lock
|
|
24
|
+
pip-selfcheck.json
|
|
25
|
+
|
|
26
|
+
# === Testing & coverage ===
|
|
27
|
+
.pytest_cache/
|
|
28
|
+
.coverage
|
|
29
|
+
.coverage.*
|
|
30
|
+
coverage.xml
|
|
31
|
+
htmlcov/
|
|
32
|
+
.tox/
|
|
33
|
+
.nox/
|
|
34
|
+
|
|
35
|
+
# === Type checkers ===
|
|
36
|
+
.mypy_cache/
|
|
37
|
+
.pyright/
|
|
38
|
+
.dmypy.json
|
|
39
|
+
pyrightconfig.json
|
|
40
|
+
|
|
41
|
+
# === PyTorch artefacts ===
|
|
42
|
+
*.pt
|
|
43
|
+
*.pth
|
|
44
|
+
*.ckpt
|
|
45
|
+
*.bin
|
|
46
|
+
*.safetensors
|
|
47
|
+
|
|
48
|
+
# === ONNX / NNEF artefacts ===
|
|
49
|
+
*.onnx
|
|
50
|
+
*.nnef
|
|
51
|
+
*.nnef.tar
|
|
52
|
+
*.rgg
|
|
53
|
+
|
|
54
|
+
# === Jupyter ===
|
|
55
|
+
.ipynb_checkpoints/
|
|
56
|
+
*.ipynb
|
|
57
|
+
|
|
58
|
+
# === IDEs ===
|
|
59
|
+
.idea/
|
|
60
|
+
.vscode/
|
|
61
|
+
*.sublime-project
|
|
62
|
+
*.sublime-workspace
|
|
63
|
+
*.code-workspace
|
|
64
|
+
*.iml
|
|
65
|
+
.metals/
|
|
66
|
+
.bloop/
|
|
67
|
+
|
|
68
|
+
# === macOS ===
|
|
69
|
+
.DS_Store
|
|
70
|
+
.AppleDouble
|
|
71
|
+
.LSOverride
|
|
72
|
+
._*
|
|
73
|
+
.Spotlight-V100
|
|
74
|
+
.Trashes
|
|
75
|
+
|
|
76
|
+
# === Linux ===
|
|
77
|
+
*~
|
|
78
|
+
.fuse_hidden*
|
|
79
|
+
.Trash-*
|
|
80
|
+
|
|
81
|
+
# === Windows ===
|
|
82
|
+
Thumbs.db
|
|
83
|
+
ehthumbs.db
|
|
84
|
+
Desktop.ini
|
|
85
|
+
|
|
86
|
+
# === Logs & temp ===
|
|
87
|
+
*.log
|
|
88
|
+
*.tmp
|
|
89
|
+
*.bak
|
|
90
|
+
*.swp
|
|
91
|
+
*.swo
|
|
92
|
+
|
|
93
|
+
# === neuropress output ===
|
|
94
|
+
/data/
|
|
95
|
+
vast_output/
|
|
96
|
+
.ruff_cache/
|
|
97
|
+
.ty_cache/
|
|
98
|
+
papers/
|
|
99
|
+
.cachesize
|
|
100
|
+
.tmp_inst.json
|
|
@@ -0,0 +1,16 @@
|
|
|
1
|
+
repos:
|
|
2
|
+
- repo: https://github.com/astral-sh/ruff-pre-commit
|
|
3
|
+
rev: v0.11.5
|
|
4
|
+
hooks:
|
|
5
|
+
- id: ruff
|
|
6
|
+
args: [--fix, --exit-non-zero-on-fix]
|
|
7
|
+
- id: ruff-format
|
|
8
|
+
|
|
9
|
+
- repo: local
|
|
10
|
+
hooks:
|
|
11
|
+
- id: ty
|
|
12
|
+
name: ty
|
|
13
|
+
entry: uv run ty check
|
|
14
|
+
language: system
|
|
15
|
+
types: [python]
|
|
16
|
+
pass_filenames: false
|
|
@@ -0,0 +1,43 @@
|
|
|
1
|
+
strictness: medium
|
|
2
|
+
doc-warnings: false
|
|
3
|
+
test-warnings: false
|
|
4
|
+
|
|
5
|
+
# ruff already handles pycodestyle and pyflakes
|
|
6
|
+
without-tools:
|
|
7
|
+
- pycodestyle
|
|
8
|
+
- pyflakes
|
|
9
|
+
|
|
10
|
+
ignore-patterns:
|
|
11
|
+
- .venv
|
|
12
|
+
|
|
13
|
+
mccabe:
|
|
14
|
+
options:
|
|
15
|
+
# GPTQ and hessian methods are inherently complex multi-branch algorithms.
|
|
16
|
+
max-complexity: 20
|
|
17
|
+
|
|
18
|
+
pylint:
|
|
19
|
+
options:
|
|
20
|
+
# Scientific / ML code naturally has many parameters and local variables.
|
|
21
|
+
max-args: 10
|
|
22
|
+
max-locals: 28
|
|
23
|
+
max-branches: 20
|
|
24
|
+
# Allow extension-pkg access (torch._C, torch.cuda, etc.)
|
|
25
|
+
extension-pkg-allow-list:
|
|
26
|
+
- torch
|
|
27
|
+
|
|
28
|
+
disable:
|
|
29
|
+
# Lazy / conditional imports are intentional in this codebase
|
|
30
|
+
# (OffloadedTensor, datasets, GranularityType, os, …).
|
|
31
|
+
- import-outside-toplevel
|
|
32
|
+
# Hook callbacks must match PyTorch signatures; unused args are structural.
|
|
33
|
+
- unused-argument
|
|
34
|
+
# Accessing _clone_weights / _LinAlgError on peer classes is intentional.
|
|
35
|
+
- protected-access
|
|
36
|
+
# torch.linalg.cholesky is callable; pylint static analysis is wrong.
|
|
37
|
+
- not-callable
|
|
38
|
+
# torch._C members are C-extension; unavailable at static analysis time.
|
|
39
|
+
- c-extension-no-member
|
|
40
|
+
# Project does not use Django.
|
|
41
|
+
- django-not-available
|
|
42
|
+
# Scientific APIs naturally have more than 5 positional parameters.
|
|
43
|
+
- too-many-positional-arguments
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
3.13
|
|
@@ -0,0 +1,110 @@
|
|
|
1
|
+
# neuropress architecture
|
|
2
|
+
|
|
3
|
+
## Commands
|
|
4
|
+
|
|
5
|
+
```bash
|
|
6
|
+
# Install dependencies (Python 3.13, managed with uv)
|
|
7
|
+
uv sync
|
|
8
|
+
|
|
9
|
+
# Install with optional LLM data loader (HuggingFace datasets)
|
|
10
|
+
uv sync --extra llm
|
|
11
|
+
|
|
12
|
+
# Run all unit tests
|
|
13
|
+
uv run pytest
|
|
14
|
+
|
|
15
|
+
# Run a single test file
|
|
16
|
+
uv run pytest tests/test_pipeline.py -v
|
|
17
|
+
|
|
18
|
+
# Run a single test by name
|
|
19
|
+
uv run pytest tests/test_pipeline.py::TestQuantizationPipeline::test_pipeline_runs -v
|
|
20
|
+
|
|
21
|
+
# Run integration tests (downloads ~500 MB model, requires RUN_INTEGRATION=1)
|
|
22
|
+
RUN_INTEGRATION=1 uv run pytest tests/integration/ -v
|
|
23
|
+
```
|
|
24
|
+
|
|
25
|
+
## Architecture
|
|
26
|
+
|
|
27
|
+
`neuropress` is a graph-aware neural network quantization framework targeting NNEF export via `torch_to_nnef`. The full pipeline proceeds in six steps (see `pipeline.py`):
|
|
28
|
+
|
|
29
|
+
1. **QGraph** – `RuntimeModuleGraphProvider` traces a single forward pass using PyTorch module hooks to build a `QGraph` (`nodes: Dict[str, QNode]`, `values: Dict[str, QValue]`). No symbolic tracing (FX/Dynamo) is used.
|
|
30
|
+
|
|
31
|
+
2. **Targets** – `WeightTarget` (`module_id/param_name`) and `ActivationTarget` (`value_id`) identify what to quantize.
|
|
32
|
+
|
|
33
|
+
3. **StatisticsRegistry** – `QuantizationParamEstimator.declare_statistics()` returns `StatisticRequest` objects. The registry deduplicates them (key = `target_key::stat_type`), installs hooks, runs a forward pass over calibration data, and stores results. Supported stat types: `"min_max"` and `"hessian"`.
|
|
34
|
+
|
|
35
|
+
4. **QuantizationRecipe** – ties together `allowed_weight_qschemes`, `allowed_activation_qschemes`, a `QuantizationParamEstimator`, and a `ModuleQuantizationAdapter`.
|
|
36
|
+
|
|
37
|
+
5. **QuantizationPlan** – the output of estimation: dicts mapping target keys to `ZpScaleScheme` or `QSchemeGroup` instances.
|
|
38
|
+
|
|
39
|
+
6. **ModuleQuantizationAdapter** – applies the plan to the model (abstract; user-supplied).
|
|
40
|
+
|
|
41
|
+
### Key abstractions
|
|
42
|
+
|
|
43
|
+
| Class | Role |
|
|
44
|
+
|---|---|
|
|
45
|
+
| `QSchemeDescriptor` | Static backend constraint (nbits, symmetric, granularity) |
|
|
46
|
+
| `ZpScaleScheme` | Concrete affine scheme (zero_point, scale, nbits); extends `torch_to_nnef.QScheme` |
|
|
47
|
+
| `QSchemeGroup` / `PanelSpec` / `TensorSlice` | Mixed-precision grouping of weight regions |
|
|
48
|
+
| `CalibrationMethod` | ABC for algorithms; subclasses implement `update_statistics` + `compute_qscheme` |
|
|
49
|
+
| `CalibrationMethodOnActivation` | Subclass of above for Hessian-based methods (GPTQ) |
|
|
50
|
+
| `GPTQCalibration` | Full GPTQ implementation (Frantar et al. ICLR 2023) |
|
|
51
|
+
| `GPTQMixedPrecisionCalibration` | GPTQ with per-column-group bit-width allocation via `MixedAllocationStrategy` |
|
|
52
|
+
| `MixedPrecisionLinear` | Drop-in replacement for `nn.Linear` produced by `split_mixed_precision_linear` |
|
|
53
|
+
|
|
54
|
+
### Calibration data shapes (4-D convention)
|
|
55
|
+
|
|
56
|
+
Weights are reshaped to `[rows, 1, n_groups, group_size]` before quantization. Odd dims (1, group_size) are "reduction" dims; even dims (rows, n_groups) are "kept". `qscheme_data_shape` is the same rank with reduction dims set to 1.
|
|
57
|
+
|
|
58
|
+
### `StatField` annotation pattern
|
|
59
|
+
|
|
60
|
+
`CalibrationMethod` subclasses declare statistic dependencies via `Annotated` type hints. The registry auto-discovers and injects them — no manual wiring needed.
|
|
61
|
+
|
|
62
|
+
```python
|
|
63
|
+
hessian: Annotated[Optional[torch.Tensor], StatField("hessian", "hessian")] = None
|
|
64
|
+
```
|
|
65
|
+
|
|
66
|
+
`StatField(stat_type, field_key)`:
|
|
67
|
+
- `stat_type` — matches the string key in `StatisticRequest` (`"hessian"`, `"hessian_diag"`, `"min_max"`).
|
|
68
|
+
- `field_key` — the attribute name to populate on the method instance.
|
|
69
|
+
|
|
70
|
+
Flow: `StatisticsRegistry.register_method(method, target_key)` scans `Annotated` hints to build `StatisticRequest` objects. After collection, `registry.inject(method, target_key)` sets the annotated fields from stored results. The pattern lives across three files: `calibration/probe/collectors.py` (StatField), `calibration/probe/registry.py` (discover + inject), and each `CalibrationMethod` subclass.
|
|
71
|
+
|
|
72
|
+
### Estimator config / implementation split
|
|
73
|
+
|
|
74
|
+
`VptqGroupEstimatorConfig` (`estimator_configs.py`) is the user-facing Pydantic model; `VPTQGroupEstimator` (`estimators.py`) is the runtime dataclass. The config's `build()` passes every field explicitly, so **`VPTQGroupEstimator` is the single source of truth for defaults** — `VptqGroupEstimatorConfig` derives its defaults from the dataclass fields via `_vd()` at class definition time. If you add a field to `VPTQGroupEstimator`, add a corresponding `_vd("field_name")` default in the config.
|
|
75
|
+
|
|
76
|
+
### Pipeline wiring: `_cache` and `_spans`
|
|
77
|
+
|
|
78
|
+
`QuantizationParamEstimator` declares `_cache: CacheManager` and `_spans: SpanCollector` as class attributes with no-op defaults. The pipeline calls `set_cache()` and `set_span_collector()` before `prepare()`. Subclasses use `self._cache` / `self._spans` directly — no `getattr` fallback needed.
|
|
79
|
+
|
|
80
|
+
### Compression registry (`neuropress/t2n/registry.py`)
|
|
81
|
+
|
|
82
|
+
Follows the `torch_to_nnef` registry pattern: a dict (`NEUROPRESS_COMPRESSION`) maps string keys to callables with signature `fn(model, tokenizer=..., **kwargs) -> model`. Pass it to `LLMExporter.prepare`:
|
|
83
|
+
|
|
84
|
+
```python
|
|
85
|
+
exporter.prepare(
|
|
86
|
+
compression_method="open_text_c4_gptq_q4_0",
|
|
87
|
+
compression_registry="neuropress.t2n.registry.NEUROPRESS_COMPRESSION",
|
|
88
|
+
)
|
|
89
|
+
```
|
|
90
|
+
|
|
91
|
+
Every entry is built by `_make_entry(config)`, which binds a `QuantizationRecipeConfig` into a `functools.partial` of `run_quantization_pipeline` and stashes the config JSON on the partial as `__config__` so the manifest can record it. Calling an entry:
|
|
92
|
+
|
|
93
|
+
1. Reads `QuantizationRecipe.calibration` and loads batches from `calibration.DATASET_REGISTRY`
|
|
94
|
+
2. Transfers batches to the model device one at a time (no OOM on large calibration sets)
|
|
95
|
+
3. Delegates to `QuantizationPipeline`, which runs the six steps above
|
|
96
|
+
4. Applies kernel lowering when the recipe configures one
|
|
97
|
+
|
|
98
|
+
Recipes can also be built from JSON via `QuantizationRecipeConfig.from_json(...).build()`, which resolves component types from the registries at deserialization time.
|
|
99
|
+
|
|
100
|
+
Layers whose `in_features` is not a multiple of the group size are skipped (Q4_0 group size = 32).
|
|
101
|
+
|
|
102
|
+
Registry entries: `"open_text_c4_gptq_q4_0"` (128 samples), `"open_text_c4_gptq_q4_0_32samples"`, `"open_text_c4_gptq_q4_0_256samples"`.
|
|
103
|
+
|
|
104
|
+
### LLM calibration data
|
|
105
|
+
|
|
106
|
+
`neuropress.calibration.data.llm.load_c4_calibration(tokenizer, ...)` loads C4 samples as `{"input_ids": Tensor[batch_size, seq_len]}` dicts. Requires the `llm` optional dependency (`datasets`).
|
|
107
|
+
|
|
108
|
+
### Integration test pattern
|
|
109
|
+
|
|
110
|
+
The integration test (`tests/integration/test_gptq_nnef_export.py`) demonstrates the end-to-end workflow: HuggingFace model load → C4 calibration → layer-by-layer GPTQ → `torch_to_nnef` NNEF export. It is gated behind `RUN_INTEGRATION=1` and also requires `transformers` and `datasets`.
|
|
@@ -0,0 +1,82 @@
|
|
|
1
|
+
# Dockerfile.vast — pre-baked base image for neuropress vast.ai GPU jobs.
|
|
2
|
+
#
|
|
3
|
+
# Cuts cold-start time from ~15 min to ~3 min by pre-warming the uv package
|
|
4
|
+
# cache inside the image layer. On a fresh instance, uv sync only needs to
|
|
5
|
+
# install the neuropress source itself (a few seconds), not re-download torch
|
|
6
|
+
# or any other dependency.
|
|
7
|
+
#
|
|
8
|
+
# Build and push once:
|
|
9
|
+
#
|
|
10
|
+
# DOCKER_BUILDKIT=1 docker build -f Dockerfile.vast -t <user>/neuropress-base:latest .
|
|
11
|
+
# docker push <user>/neuropress-base:latest
|
|
12
|
+
#
|
|
13
|
+
# Subsequent rebuilds after a pyproject.toml change are fast (~30 s) because
|
|
14
|
+
# BuildKit caches the downloaded packages on the build host. The first build
|
|
15
|
+
# still downloads everything (~9 min).
|
|
16
|
+
#
|
|
17
|
+
# Then tell neuropress-vast to use it:
|
|
18
|
+
#
|
|
19
|
+
# export NEUROPRESS_DOCKER_IMAGE=<user>/neuropress-base:latest
|
|
20
|
+
# neuropress-vast --job neuropress-eval-compare ...
|
|
21
|
+
#
|
|
22
|
+
# The image must be public (vast.ai pulls without auth by default).
|
|
23
|
+
# Re-build any time pyproject.toml dependencies change.
|
|
24
|
+
|
|
25
|
+
# Slim CUDA 12.4 runtime — ~800 MB compressed vs ~5 GB for pytorch/pytorch.
|
|
26
|
+
# PyTorch is installed via uv from PyPI below, so we don't need it in the base.
|
|
27
|
+
FROM nvidia/cuda:13.0.2-cudnn-runtime-ubuntu24.04
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
ENV DEBIAN_FRONTEND=noninteractive
|
|
31
|
+
|
|
32
|
+
# System tools needed by setup commands:
|
|
33
|
+
# tmux — disconnect-resilient job executor
|
|
34
|
+
# build-essential — g++/make for torch-to-nnef C++ extension
|
|
35
|
+
# ninja-build — required by torch.utils.cpp_extension for CUDA JIT kernels
|
|
36
|
+
# cuda-nvcc — nvcc compiler for CUDA JIT (runtime images lack it)
|
|
37
|
+
# curl — uv installer
|
|
38
|
+
# git — some HuggingFace libs call git
|
|
39
|
+
#
|
|
40
|
+
# Cache mounts keep downloaded .deb files on the build host so subsequent
|
|
41
|
+
# builds skip re-downloading them. The mounts are not part of the image
|
|
42
|
+
# layer, so no rm -rf /var/lib/apt/lists/* is needed.
|
|
43
|
+
RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
|
|
44
|
+
--mount=type=cache,target=/var/lib/apt,sharing=locked \
|
|
45
|
+
apt-get update -qq && \
|
|
46
|
+
apt-get install -y --no-install-recommends \
|
|
47
|
+
tmux build-essential ninja-build cuda-nvcc curl ca-certificates git
|
|
48
|
+
|
|
49
|
+
# Install uv
|
|
50
|
+
RUN curl -LsSf https://astral.sh/uv/install.sh | sh
|
|
51
|
+
ENV PATH="/root/.local/bin:$PATH"
|
|
52
|
+
ENV UV_CACHE_DIR=/root/.cache/uv
|
|
53
|
+
|
|
54
|
+
# Pre-install Python 3.13.
|
|
55
|
+
# Cache mount keeps the CPython tarball on the build host; the installed
|
|
56
|
+
# binary goes to ~/.local/share/uv/python (baked into the image layer).
|
|
57
|
+
RUN --mount=type=cache,target=/root/.cache/uv \
|
|
58
|
+
uv python install 3.13
|
|
59
|
+
|
|
60
|
+
# ── Pre-warm the uv package cache ────────────────────────────────────────────
|
|
61
|
+
# Copy only the dependency manifest (not the source) so this layer is rebuilt
|
|
62
|
+
# only when dependencies change, not on every code change.
|
|
63
|
+
WORKDIR /tmp/neuropress-seed
|
|
64
|
+
COPY pyproject.toml ./
|
|
65
|
+
# hatchling needs a minimal package layout to resolve the project
|
|
66
|
+
RUN mkdir -p neuropress && \
|
|
67
|
+
printf '"""neuropress."""\n' > neuropress/__init__.py && \
|
|
68
|
+
touch README.md
|
|
69
|
+
|
|
70
|
+
# Download packages to the BuildKit cache (/tmp/uv-cache, preserved on the
|
|
71
|
+
# build host between runs). On first build this takes ~9 min; on subsequent
|
|
72
|
+
# rebuilds after a dependency change uv resolves from cache in ~30 s.
|
|
73
|
+
# After sync, cp bakes the populated cache into UV_CACHE_DIR (image layer) so
|
|
74
|
+
# that `uv sync` on a fresh vast.ai instance needs no network downloads.
|
|
75
|
+
RUN --mount=type=cache,target=/tmp/uv-cache,sharing=locked \
|
|
76
|
+
UV_CACHE_DIR=/tmp/uv-cache \
|
|
77
|
+
uv sync --extra llm --extra eval --extra dev && \
|
|
78
|
+
rm -rf .venv && \
|
|
79
|
+
mkdir -p "$UV_CACHE_DIR" && \
|
|
80
|
+
cp -a /tmp/uv-cache/. "$UV_CACHE_DIR/"
|
|
81
|
+
|
|
82
|
+
WORKDIR /workspace
|