From cec125bc6972a2d4437f9edc66c4e418399ae860 Mon Sep 17 00:00:00 2001 From: Reiase Date: Fri, 24 Jul 2026 18:11:42 +0800 Subject: [PATCH] Enhance CI/CD workflows and documentation for nightly builds - Introduced a new GitHub Action for setting up the build environment, including Python, Rust, and optional protoc installation. - Added CI workflows for linting, testing, and nightly builds, ensuring consistent checks across platforms. - Implemented a script for installing the latest nightly wheel from GitHub Releases, improving accessibility for users. - Updated installation documentation to include instructions for nightly builds and usage examples. - Refined existing documentation for clarity and completeness. These changes streamline the development process and enhance the usability of nightly builds for users. --- .github/actions/setup-build-env/action.yml | 58 +++++++ .github/workflows/ci.yml | 86 ++++++++++ .github/workflows/nightly.yml | 156 ++++++++++++++++++ .../src/proxy/network_policy.rs | 2 +- docs/src/installation.md | 8 + docs/src/installation.zh.md | 8 + examples/03_sampler.py | 12 +- examples/06_tiered_lazy_prefetch_eviction.py | 4 + examples/capture-walkthrough/mock_llm.py | 5 +- justfile | 118 ++++++++++++- persisting/__init__.py | 17 +- persisting/core.py | 11 +- persisting/queue/backend.py | 11 +- persisting/queue/metadata.py | 4 +- persisting/queue/queue.py | 18 +- persisting/queue/status_tracker.py | 4 +- persisting/sampler/reader.py | 2 +- persisting/search/__init__.py | 8 +- persisting/trajectory.py | 4 +- scripts/ci/set_nightly_local_version.py | 32 ++++ scripts/install-nightly.sh | 89 ++++++++++ tests/test_block_io.py | 1 + tests/test_block_store.py | 24 ++- tests/test_block_store_actor.py | 6 + tests/test_local_tensor_store.py | 15 +- tests/test_metadata.py | 12 +- tests/test_queue_tensordict.py | 4 +- tests/test_taa_subscript_planning.py | 2 + 28 files changed, 654 insertions(+), 67 deletions(-) create mode 100644 .github/actions/setup-build-env/action.yml create mode 100644 .github/workflows/ci.yml create mode 100644 .github/workflows/nightly.yml create mode 100755 scripts/ci/set_nightly_local_version.py create mode 100755 scripts/install-nightly.sh diff --git a/.github/actions/setup-build-env/action.yml b/.github/actions/setup-build-env/action.yml new file mode 100644 index 0000000..ee22a60 --- /dev/null +++ b/.github/actions/setup-build-env/action.yml @@ -0,0 +1,58 @@ +name: "Setup Build Environment" +description: "Python, Rust, uv, and optional protoc for Persisting CI" +inputs: + python-version: + description: "Python version" + required: false + default: "3.12" + rust-toolchain: + description: "Rust toolchain" + required: false + default: "stable" + install-protoc: + description: "Install protobuf compiler (needed by persisting-dlcapt)" + required: false + default: "true" + components: + description: "Rust components (comma-separated)" + required: false + default: "rustfmt, clippy" + +runs: + using: "composite" + steps: + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: ${{ inputs.python-version }} + + - name: Install Rust toolchain + uses: dtolnay/rust-toolchain@master + with: + toolchain: ${{ inputs.rust-toolchain }} + components: ${{ inputs.components }} + + - name: Install uv + uses: astral-sh/setup-uv@v6 + with: + enable-cache: true + + - name: Install protoc + if: inputs.install-protoc == 'true' && runner.os == 'Linux' + shell: bash + run: | + sudo apt-get update -qq + sudo apt-get install -y -qq --no-install-recommends protobuf-compiler + protoc --version + + - name: Install protoc + if: inputs.install-protoc == 'true' && runner.os == 'macOS' + shell: bash + run: | + brew install protobuf + protoc --version + + - name: Point PyO3 at CI Python + shell: bash + run: | + echo "PYO3_PYTHON=$(python -c 'import sys; print(sys.executable)')" >> "$GITHUB_ENV" diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..597ffea --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,86 @@ +name: CI + +on: + push: + branches: [main, master] + pull_request: + branches: [main, master] + +concurrency: + group: ci-${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +env: + CARGO_TERM_COLOR: always + RUST_BACKTRACE: 1 + +jobs: + lint: + name: Lint & Format + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Setup Build Environment + uses: ./.github/actions/setup-build-env + with: + install-protoc: "true" + + - uses: Swatinem/rust-cache@v2 + + - name: Rust formatting + run: cargo fmt --all -- --check + + - name: Rust clippy + # Match local `just clippy` for now; tighten to `-D warnings` once the + # workspace is clean (same bar as Pulsing). + run: cargo clippy --workspace --all-targets --locked + + - name: Python lint (ruff) + run: uvx ruff check persisting/ + + rust-test: + name: Rust Test (${{ matrix.os }}) + runs-on: ${{ matrix.os }} + strategy: + fail-fast: false + matrix: + os: [ubuntu-latest, macos-latest] + steps: + - uses: actions/checkout@v4 + + - name: Setup Build Environment + uses: ./.github/actions/setup-build-env + with: + install-protoc: "true" + + - uses: Swatinem/rust-cache@v2 + + - name: Build + run: cargo build --workspace --all-targets --locked + + - name: Test + run: cargo test --workspace --locked + + python-test: + name: Python Test + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Setup Build Environment + uses: ./.github/actions/setup-build-env + with: + python-version: "3.12" + install-protoc: "true" + + - uses: Swatinem/rust-cache@v2 + + - name: Sync Python deps + run: uv sync --extra all --extra dev + + - name: Build extension (maturin develop) + run: uv run maturin develop + + - name: Pytest + run: uv run pytest tests/ -q diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml new file mode 100644 index 0000000..be7288e --- /dev/null +++ b/.github/workflows/nightly.yml @@ -0,0 +1,156 @@ +# Nightly wheels from main — rolling GitHub Release tag `nightly`. +# +# Install (after a successful run): +# curl -fsSL https://raw.githubusercontent.com/DeepLink-org/Persisting/main/scripts/install-nightly.sh | bash +# Or browse assets: https://github.com/DeepLink-org/Persisting/releases/tag/nightly + +name: Nightly Build + +on: + schedule: + # 03:00 UTC daily + - cron: "0 3 * * *" + workflow_dispatch: + push: + branches: [main, master] + +permissions: + contents: read + +concurrency: + group: nightly-${{ github.ref }} + cancel-in-progress: true + +jobs: + meta: + if: github.repository == 'DeepLink-org/Persisting' + runs-on: ubuntu-latest + outputs: + local_version: ${{ steps.ver.outputs.local_version }} + steps: + - id: ver + shell: bash + run: | + short_sha="${GITHUB_SHA:0:7}" + echo "local_version=g${GITHUB_RUN_NUMBER}.${short_sha}" >> "$GITHUB_OUTPUT" + + linux: + if: github.repository == 'DeepLink-org/Persisting' + needs: meta + runs-on: ubuntu-latest + strategy: + matrix: + target: [x86_64] + steps: + - uses: actions/checkout@v4 + + - name: Setup Build Environment + uses: ./.github/actions/setup-build-env + with: + python-version: "3.12" + install-protoc: "true" + + - uses: Swatinem/rust-cache@v2 + + - name: Apply nightly local version (PEP 440) + run: python scripts/ci/set_nightly_local_version.py "${{ needs.meta.outputs.local_version }}" + + - name: Build wheels + uses: PyO3/maturin-action@v1 + with: + target: ${{ matrix.target }} + args: >- + --release --out dist --zig --compatibility manylinux2014 + --auditwheel repair + sccache: "true" + manylinux: auto + rust-toolchain: stable + + - name: Upload wheels + uses: actions/upload-artifact@v4 + with: + name: nightly-wheels-linux-${{ matrix.target }} + path: dist/*.whl + retention-days: 14 + + macos: + if: github.repository == 'DeepLink-org/Persisting' + needs: meta + runs-on: macos-latest + strategy: + matrix: + target: [x86_64, aarch64] + steps: + - uses: actions/checkout@v4 + + - name: Setup Build Environment + uses: ./.github/actions/setup-build-env + with: + python-version: "3.12" + install-protoc: "true" + + - uses: Swatinem/rust-cache@v2 + + - name: Apply nightly local version (PEP 440) + run: python scripts/ci/set_nightly_local_version.py "${{ needs.meta.outputs.local_version }}" + + - name: Build wheels + uses: PyO3/maturin-action@v1 + with: + target: ${{ matrix.target }} + args: >- + --release --out dist --find-interpreter + sccache: "true" + rust-toolchain: stable + + - name: Upload wheels + uses: actions/upload-artifact@v4 + with: + name: nightly-wheels-macos-${{ matrix.target }} + path: dist/*.whl + retention-days: 14 + + publish: + name: Publish nightly release + if: github.repository == 'DeepLink-org/Persisting' + needs: [meta, linux, macos] + runs-on: ubuntu-latest + permissions: + contents: write + steps: + - uses: actions/checkout@v4 + + - name: Download wheels + uses: actions/download-artifact@v4 + with: + pattern: nightly-wheels-* + merge-multiple: true + path: dist + + - name: List artifacts + run: ls -la dist/ + + - name: Publish rolling nightly release + uses: ncipollo/release-action@v1 + with: + tag: nightly + name: Nightly (main) + commit: ${{ github.sha }} + allowUpdates: true + makeLatest: false + prerelease: true + removeArtifacts: true + artifactErrorsFailBuild: true + artifacts: dist/*.whl + body: | + Rolling nightly wheels built from [`${{ github.sha }}`](${{ github.server_url }}/${{ github.repository }}/commit/${{ github.sha }}) on **${{ github.ref_name }}**. + + **Quick install** (auto-detect platform / Python): + + ```bash + curl -fsSL https://raw.githubusercontent.com/${{ github.repository }}/main/scripts/install-nightly.sh | bash + ``` + + Or pick a `.whl` below and `pip install --force-reinstall `. + + Built with local version `+${{ needs.meta.outputs.local_version }}` (base version from `pyproject.toml`). diff --git a/crates/persisting-capture/src/proxy/network_policy.rs b/crates/persisting-capture/src/proxy/network_policy.rs index ac69cb1..a5fb55b 100644 --- a/crates/persisting-capture/src/proxy/network_policy.rs +++ b/crates/persisting-capture/src/proxy/network_policy.rs @@ -561,6 +561,6 @@ upstream = "http://127.0.0.1:9/v1" assert_eq!(cfg.network.allowed_hosts.len(), 3); let p = NetworkPolicy::from_config(&cfg).unwrap(); assert!(assert_egress_allowed(&p, "a.example.com").is_ok()); - assert!(!assert_egress_allowed(&p, "example.com").is_ok()); + assert!(assert_egress_allowed(&p, "example.com").is_err()); } } diff --git a/docs/src/installation.md b/docs/src/installation.md index d50d35e..14a1ee4 100644 --- a/docs/src/installation.md +++ b/docs/src/installation.md @@ -23,6 +23,14 @@ cd Persisting pip install -e ".[lance]" ``` +### Nightly (pre-release wheels) + +Daily / `main` pushes refresh the rolling GitHub Release tag [`nightly`](https://github.com/DeepLink-org/Persisting/releases/tag/nightly): + +```bash +curl -fsSL https://raw.githubusercontent.com/DeepLink-org/Persisting/main/scripts/install-nightly.sh | bash +``` + ## Verify ```python diff --git a/docs/src/installation.zh.md b/docs/src/installation.zh.md index be6fca9..f135816 100644 --- a/docs/src/installation.zh.md +++ b/docs/src/installation.zh.md @@ -23,6 +23,14 @@ cd Persisting pip install -e ".[lance]" ``` +### Nightly（预发布 wheel） + +每日 / `main` 推送会滚动更新 GitHub Release 标签 [`nightly`](https://github.com/DeepLink-org/Persisting/releases/tag/nightly)： + +```bash +curl -fsSL https://raw.githubusercontent.com/DeepLink-org/Persisting/main/scripts/install-nightly.sh | bash +``` + ## 验证 ```python diff --git a/examples/03_sampler.py b/examples/03_sampler.py index a080685..402651e 100644 --- a/examples/03_sampler.py +++ b/examples/03_sampler.py @@ -41,7 +41,9 @@ def sample( *args: Any, **kwargs: Any, ) -> tuple[list[int], list[int]]: - taken = ready_indexes[-batch_size:] if len(ready_indexes) >= batch_size else ready_indexes[:] + taken = ( + ready_indexes[-batch_size:] if len(ready_indexes) >= batch_size else ready_indexes[:] + ) taken = list(reversed(taken)) return taken, taken @@ -64,7 +66,9 @@ async def demo_sequential_sampler() -> None: print(f"\n Reading in batches of {batch_size} with SequentialSampler:") for step in range(3): - records, offset = await queue.get_batch(sampler, batch_size=batch_size, consumed_offset=offset) + records, offset = await queue.get_batch( + sampler, batch_size=batch_size, consumed_offset=offset + ) if not records: break print(f" Step {step + 1}: got {len(records)} records, new offset = {offset}") @@ -91,7 +95,9 @@ async def demo_custom_reverse_sampler() -> None: print("\n Reading in batches of 2 with ReverseSampler (last indices first):") for step in range(3): - records, offset = await queue.get_batch(sampler, batch_size=batch_size, consumed_offset=offset) + records, offset = await queue.get_batch( + sampler, batch_size=batch_size, consumed_offset=offset + ) if not records: break print(f" Step {step + 1}: got {len(records)} records") diff --git a/examples/06_tiered_lazy_prefetch_eviction.py b/examples/06_tiered_lazy_prefetch_eviction.py index b9f777d..ab7f395 100644 --- a/examples/06_tiered_lazy_prefetch_eviction.py +++ b/examples/06_tiered_lazy_prefetch_eviction.py @@ -56,6 +56,7 @@ # ── Prefill ──────────────────────────────────────────────── + def prefill(sid: str, seq_len: int): for layer in range(NUM_LAYERS): for head in range(NUM_HEADS): @@ -68,6 +69,7 @@ def prefill(sid: str, seq_len: int): # prefetch 需要空间 → 按选取策略挑 unpinned block 驱逐 # pin 保护 → unpin 后回到候选池 + def decode_step(sid: str, pos: int, stream): for layer in range(NUM_LAYERS): if layer + 1 < NUM_LAYERS: @@ -86,6 +88,7 @@ def decode_step(sid: str, pos: int, stream): # 这比被动 LRU 更快腾出空间（LRU 可能选错——刚被抢占的 session # 在 LRU 里还是"热"的，但调度器知道它不会再用了）。 + def serve_batch(batch: list[tuple[str, int]], preempted: list[str], stream): for sid in preempted: kv[sid, :, :, :].evict() @@ -102,6 +105,7 @@ def serve_batch(batch: list[tuple[str, int]], preempted: list[str], stream): # ── Main ────────────────────────────────────────────────── + def main(): stream = cuda.Stream() diff --git a/examples/capture-walkthrough/mock_llm.py b/examples/capture-walkthrough/mock_llm.py index 059e449..795f4d2 100755 --- a/examples/capture-walkthrough/mock_llm.py +++ b/examples/capture-walkthrough/mock_llm.py @@ -100,8 +100,5 @@ def log_message(self, fmt: str, *args: object) -> None: if __name__ == "__main__": - print( - f"Mock LLM: http://{HOST}:{PORT}" - f" (/v1/chat/completions, /v1/messages) Ctrl+C 停止" - ) + print(f"Mock LLM: http://{HOST}:{PORT} (/v1/chat/completions, /v1/messages) Ctrl+C 停止") HTTPServer((HOST, PORT), Handler).serve_forever() diff --git a/justfile b/justfile index da42dcf..5215756 100644 --- a/justfile +++ b/justfile @@ -10,6 +10,11 @@ docs_dir := repo / "docs" gen_py := repo / "scripts" / "generate_benchmark_data.py" test_suite_sh := repo / "scripts" / "test_suite.sh" +# Python 路径（ruff format） +ruff_paths := "persisting tests examples" +# lint 默认只扫包代码（与 CI 一致）；全量用 lint-py-all +ruff_lint_paths := "persisting" + engine_filename := if os() == "macos" { "libpersisting_engine.dylib" } else if os() == "windows" { @@ -24,11 +29,17 @@ default: @just --list --unsorted @echo "" @echo "常用：" + @echo " just fmt # 格式化 Rust + Python" + @echo " just lint # clippy + ruff check" + @echo " just fix # 自动 format + ruff --fix" @echo " just test-suite # 测试 / 回归选单（推荐入口）" @echo " just test-suite list # 列出全部套件" @echo " just dev # 日常开发（fmt + clippy + check-quick）" - @echo " just gate # 提交前（fmt + clippy + test-rust）" + @echo " just gate # 提交前（fmt + lint + test-rust）" + @echo " just ci-lint # CI lint（只检查、不改写）" @echo " just ci # CI 近似全量" + @echo " just py-dev # maturin develop（Python 扩展）" + @echo " just build-wheel # 打 release wheel → dist/" @echo " just capture-all # 全部 capture 集成" @echo " just docs-serve # 本地文档" @@ -62,27 +73,95 @@ build profile="debug": build-release: just build profile=release +# maturin release wheel（当前平台）→ dist/ +build-wheel: + #!/usr/bin/env bash + set -euo pipefail + mkdir -p dist + uvx maturin build --release -o dist + ls -la dist/*.whl + +# 开发调试 wheel（不 strip） +build-wheel-debug: + #!/usr/bin/env bash + set -euo pipefail + mkdir -p dist + uvx maturin build -o dist + ls -la dist/*.whl + clean: cargo clean + rm -rf dist target/wheels .venv htmlcov .coverage coverage.xml -# ── Rust 质量 ──────────────────────────────────────────────────────────────── +# ── 格式化 / Lint ───────────────────────────────────────────────────────────── -fmt: +# 格式化 Rust + Python（会改写文件） +fmt: fmt-rust fmt-py + +fmt-rust: cargo fmt --all +fmt-py: + uvx ruff format {{ ruff_paths }} + +# 只检查格式，不改写（CI / pre-commit） +fmt-check: fmt-check-rust fmt-check-py + +fmt-check-rust: + cargo fmt --all -- --check + +fmt-check-py: + uvx ruff format --check {{ ruff_paths }} + +# clippy + ruff（不改写） +lint: lint-rust lint-py + +lint-rust: + cargo clippy --workspace --all-targets --locked + +lint-py: + uvx ruff check {{ ruff_lint_paths }} + +# 含 tests/examples（较严，可能有存量告警） +lint-py-all: + uvx ruff check {{ ruff_paths }} + +# 与 Pulsing 同级的严格 clippy（workspace 未清干净前慎用） +clippy-deny: + cargo clippy --workspace --all-targets --locked -- -D warnings + +# 兼容旧名 clippy: - cargo clippy --workspace --all-targets + just lint-rust + +# 自动修：format + ruff --fix +fix: fmt + uvx ruff check {{ ruff_paths }} --fix + +# 仅修 Python +fix-py: fmt-py + uvx ruff check {{ ruff_paths }} --fix + +# 格式 + lint 快检（不跑测试；对应 Pulsing 的 check-quick 语义） +style: fmt-check lint + @echo "✅ format + lint OK" -# fmt + clippy + 全 workspace cargo test +# fmt + lint + Rust 测试（提交前） gate: just fmt - just clippy + just lint just test-rust +# 与 GitHub Actions `ci.yml` lint 对齐（只检查、不改写） +ci-lint: + just fmt-check-rust + just lint-rust + uvx ruff check persisting/ + # 日常开发快捷路径 dev: just fmt - just clippy + just lint-rust just check-quick # CI 近似：gate + 构建 + capture 集成 + Claude 回归 @@ -94,7 +173,7 @@ ci: # ── Rust 测试 ───────────────────────────────────────────────────────────────── -# 单 crate：engine | proto | core | capture | cli +# 单 crate：engine | proto | core | capture | cli | compute | dlcapt test-crate crate: #!/usr/bin/env bash set -euo pipefail @@ -104,7 +183,9 @@ test-crate crate: core) cargo test -p persisting-core ;; capture) cargo test -p persisting-capture ;; cli) cargo test -p persisting-cli ;; - *) echo "unknown crate: {{ crate }} (engine|proto|core|capture|cli)" >&2; exit 2 ;; + compute) cargo test -p persisting-compute ;; + dlcapt) cargo test -p persisting-dlcapt ;; + *) echo "unknown crate: {{ crate }} (engine|proto|core|capture|cli|compute|dlcapt)" >&2; exit 2 ;; esac test-rust: @@ -116,23 +197,42 @@ test-capture-claude: test-capture-fixtures: cargo test -p persisting-capture --test llm_fixtures --test ag_fixture_tests +test-capture-network: + cargo test -p persisting-capture --lib network_policy + cargo test -p persisting-capture --test network_policy_http + test-engine-integration: cargo test -p persisting-engine --test search_integration _test-engine: cargo test -p persisting-engine --lib +# Rust + Python +test: test-rust test-py + # ── Python ─────────────────────────────────────────────────────────────────── py-sync: uv sync --all-extras +# 调试迭代（更快） py-dev: + uv run maturin develop + +# 接近发布形态 +py-dev-release: uv run maturin develop --release test-py: uv run pytest tests/ -q +test-py-v: + uv run pytest tests/ -v + +# 安装本地 nightly 脚本自检（需已有 GitHub nightly release） +install-nightly: + bash "{{ repo }}/scripts/install-nightly.sh" + # ── 文档（docs/ 子项目）────────────────────────────────────────────────────── docs-sync: diff --git a/persisting/__init__.py b/persisting/__init__.py index da71926..436c745 100644 --- a/persisting/__init__.py +++ b/persisting/__init__.py @@ -10,14 +10,6 @@ __version__ = "0.1.0" -from persisting.store import ( - BlockId, - BlockStore, - Handler, - LocalTensorStore, - TensorNamespace, - region_to_index, -) from persisting.queue import ( BatchMeta, KVInterface, @@ -34,6 +26,14 @@ SequentialSampler, get_sampled_batch, ) +from persisting.store import ( + BlockId, + BlockStore, + Handler, + LocalTensorStore, + TensorNamespace, + region_to_index, +) __all__ = [ "Queue", @@ -86,6 +86,7 @@ def open( """ if dtype is None: import numpy as np + dtype = np.float32 if shape is None: raise ValueError("open() 须提供 shape") diff --git a/persisting/core.py b/persisting/core.py index 8e034c5..374d6b6 100644 --- a/persisting/core.py +++ b/persisting/core.py @@ -5,19 +5,18 @@ """ from persisting._core import ( + Address, Dimension, Point, Range, - SetC, - Address, Region, + SetC, + block_read, + block_write, canonicalize, - project_prefix, is_point_query, is_range_query, - block_read, - block_write, - TieredLoop, + project_prefix, ) try: diff --git a/persisting/queue/backend.py b/persisting/queue/backend.py index faf1cdf..0fa2daa 100644 --- a/persisting/queue/backend.py +++ b/persisting/queue/backend.py @@ -15,9 +15,9 @@ import asyncio import logging +import pickle import shutil import time -import pickle from pathlib import Path from typing import Any, AsyncIterator @@ -64,7 +64,10 @@ def _infer_arrow_table(records: list[dict[str, Any]]) -> "pa.Table": ) else: arrays[name] = pa.array( - [pickle.dumps(v, protocol=pickle.HIGHEST_PROTOCOL) if v is not None else None for v in values], + [ + pickle.dumps(v, protocol=pickle.HIGHEST_PROTOCOL) if v is not None else None + for v in values + ], type=pa.binary(), ) return pa.table(arrays) @@ -240,7 +243,9 @@ async def get_data(self, batch_meta: BatchMeta, fields: list[str] | None = None) rows = await self.get_by_indices(batch_meta.global_indexes) for row in rows: for value in row.values(): - if isinstance(value, (bytes, bytearray)) and bytes(value).startswith(ZEROCOPY_PREFIX): + if isinstance(value, (bytes, bytearray)) and bytes(value).startswith( + ZEROCOPY_PREFIX + ): self._zerocopy_stats["read_envelopes"] += 1 return decode_rows(rows, fields=fields or batch_meta.field_names) diff --git a/persisting/queue/metadata.py b/persisting/queue/metadata.py index 380456f..8381e4b 100644 --- a/persisting/queue/metadata.py +++ b/persisting/queue/metadata.py @@ -84,7 +84,9 @@ def field_names(self) -> list[str]: def select_fields(self, names: list[str]) -> "BatchMeta": return BatchMeta( samples=[sample.select_fields(names) for sample in self.samples], - custom_meta={idx: meta for idx, meta in self.custom_meta.items() if idx in self.global_indexes}, + custom_meta={ + idx: meta for idx, meta in self.custom_meta.items() if idx in self.global_indexes + }, extra_info=dict(self.extra_info), ) diff --git a/persisting/queue/queue.py b/persisting/queue/queue.py index 3462337..cd0bd1c 100644 --- a/persisting/queue/queue.py +++ b/persisting/queue/queue.py @@ -153,13 +153,19 @@ async def put( ) -> BatchMeta | None: del fields if self._distributed: - if isinstance(data, dict) or (isinstance(data, list) and data and isinstance(data[0], dict)): + if isinstance(data, dict) or ( + isinstance(data, list) and data and isinstance(data[0], dict) + ): writer = await self._ensure_writer() await writer.put(data) return None bucket = await self._resolve_partition_bucket(partition_id) result = await bucket.put_tensor(data, partition_id=partition_id) - return BatchMeta.from_dict(result) if isinstance(result, dict) and "samples" in result else None + return ( + BatchMeta.from_dict(result) + if isinstance(result, dict) and "samples" in result + else None + ) if isinstance(data, dict): await self._backend.put(data) @@ -255,7 +261,9 @@ async def get_batch( return [] return await self.get_data(meta, partition_id=partition_id) - async def mark_consumed(self, task_name: str, global_indexes: list[int], partition_id: str = "default") -> None: + async def mark_consumed( + self, task_name: str, global_indexes: list[int], partition_id: str = "default" + ) -> None: if self._distributed: bucket = await self._resolve_partition_bucket(partition_id) await bucket.mark_consumed(task_name, global_indexes) @@ -289,7 +297,9 @@ async def stream( if records: yield records return - async for batch in self._backend.get_stream(limit=limit, offset=offset, wait=wait, timeout=timeout): + async for batch in self._backend.get_stream( + limit=limit, offset=offset, wait=wait, timeout=timeout + ): yield batch async def stats(self) -> dict[str, Any]: diff --git a/persisting/queue/status_tracker.py b/persisting/queue/status_tracker.py index 18ac20d..3de6d6b 100644 --- a/persisting/queue/status_tracker.py +++ b/persisting/queue/status_tracker.py @@ -28,7 +28,9 @@ def mark_produced( shapes = shapes or {} for idx in global_indexes: self.production_status[idx] = {field: "ready" for field in fields} - self.field_meta[idx] = {field: (dtypes.get(field), shapes.get(field)) for field in fields} + self.field_meta[idx] = { + field: (dtypes.get(field), shapes.get(field)) for field in fields + } def scan_ready(self, fields: list[str], task_name: str) -> list[int]: consumed = self.consumption_status.setdefault(task_name, set()) diff --git a/persisting/sampler/reader.py b/persisting/sampler/reader.py index c5c82e4..b9ac9ca 100644 --- a/persisting/sampler/reader.py +++ b/persisting/sampler/reader.py @@ -4,8 +4,8 @@ from typing import Any, Protocol, runtime_checkable -from persisting.sampler.base import BaseSampler from persisting.queue.metadata import BatchMeta +from persisting.sampler.base import BaseSampler @runtime_checkable diff --git a/persisting/search/__init__.py b/persisting/search/__init__.py index 10e0e44..fdeeb68 100644 --- a/persisting/search/__init__.py +++ b/persisting/search/__init__.py @@ -37,9 +37,7 @@ def add_documents_batch( 返回聚合结果：``added`` 为各批成功写入条数之和；``embedding_preview`` 来自首批。 """ rows = list(documents) - return _core.search_add_batch( - dataset, rows, embedding_dim=embedding_dim, chunk_size=chunk_size - ) + return _core.search_add_batch(dataset, rows, embedding_dim=embedding_dim, chunk_size=chunk_size) def query( @@ -88,9 +86,7 @@ def rebuild_indices( merge_num_indices: int | None = None, ) -> dict[str, Any]: """Call Lance `optimize_indices` (`SearchIndexRebuildRequest`).""" - return _core.search_index_rebuild( - dataset, index_name, retrain, merge_num_indices - ) + return _core.search_index_rebuild(dataset, index_name, retrain, merge_num_indices) def create_index( diff --git a/persisting/trajectory.py b/persisting/trajectory.py index 171507c..699f072 100644 --- a/persisting/trajectory.py +++ b/persisting/trajectory.py @@ -54,9 +54,7 @@ def replay( root_session_id: str | None = None, ) -> dict[str, Any]: """Page through one session (``seq`` order for Lance; block order for markdown).""" - return _core.trajectory_replay( - storage, agent_id, session_id, offset, limit, root_session_id - ) + return _core.trajectory_replay(storage, agent_id, session_id, offset, limit, root_session_id) def stats( diff --git a/scripts/ci/set_nightly_local_version.py b/scripts/ci/set_nightly_local_version.py new file mode 100755 index 0000000..f437348 --- /dev/null +++ b/scripts/ci/set_nightly_local_version.py @@ -0,0 +1,32 @@ +#!/usr/bin/env python3 +"""Append a PEP 440 local version label for nightly wheel builds.""" + +from __future__ import annotations + +import re +import sys +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[2] +VERSION_RE = re.compile(r'^(version = ")([^"+]+)(")', re.MULTILINE) + + +def patch(path: Path, local: str) -> None: + text = path.read_text(encoding="utf-8") + new, count = VERSION_RE.subn(rf"\g<1>\g<2>+{local}\g<3>", text, count=1) + if count != 1: + raise SystemExit(f"failed to patch version in {path}") + path.write_text(new, encoding="utf-8") + + +def main() -> None: + if len(sys.argv) != 2 or not sys.argv[1].strip(): + raise SystemExit("usage: set_nightly_local_version.py ") + local = sys.argv[1].strip() + patch(ROOT / "pyproject.toml", local) + patch(ROOT / "Cargo.toml", local) + print(f"patched nightly local version +{local}") + + +if __name__ == "__main__": + main() diff --git a/scripts/install-nightly.sh b/scripts/install-nightly.sh new file mode 100755 index 0000000..ed8a073 --- /dev/null +++ b/scripts/install-nightly.sh @@ -0,0 +1,89 @@ +#!/usr/bin/env bash +# Install the latest persisting nightly wheel from GitHub Releases (tag: nightly). +# +# Usage: +# curl -fsSL https://raw.githubusercontent.com/DeepLink-org/Persisting/main/scripts/install-nightly.sh | bash +# PERSISTING_GITHUB_REPO=DeepLink-org/Persisting PERSISTING_NIGHTLY_TAG=nightly ./scripts/install-nightly.sh +# +# Environment: +# PYTHON Python interpreter (default: python3) +# PERSISTING_GITHUB_REPO owner/repo (default: DeepLink-org/Persisting) +# PERSISTING_NIGHTLY_TAG release tag (default: nightly) + +set -euo pipefail + +REPO="${PERSISTING_GITHUB_REPO:-DeepLink-org/Persisting}" +TAG="${PERSISTING_NIGHTLY_TAG:-nightly}" +PYTHON="${PYTHON:-python3}" + +if ! command -v "$PYTHON" >/dev/null 2>&1; then + echo "error: Python not found ($PYTHON)" >&2 + exit 1 +fi + +case "$(uname -s)-$(uname -m)" in + Linux-x86_64) platform_re='manylinux.*x86_64' ;; + Linux-aarch64) platform_re='manylinux.*aarch64' ;; + Darwin-arm64) platform_re='macosx.*arm64' ;; + Darwin-x86_64) platform_re='macosx.*x86_64' ;; + *) + echo "error: unsupported platform $(uname -s)-$(uname -m)" >&2 + exit 1 + ;; +esac + +min_py="$("$PYTHON" -c 'import sys; print(sys.version_info.major * 10 + sys.version_info.minor)')" +if [ "$min_py" -lt 310 ]; then + echo "error: persisting wheels require Python 3.10+ (abi3-py310); got $("$PYTHON" --version)" >&2 + exit 1 +fi +api="https://api.github.com/repos/${REPO}/releases/tags/${TAG}" + +echo "Fetching nightly release assets from ${REPO} (tag=${TAG})..." >&2 + +url="$("$PYTHON" - <&2 +"$PYTHON" -m pip install --upgrade pip +"$PYTHON" -m pip install --force-reinstall "$url" +"$PYTHON" -c "import persisting; print('persisting', persisting.__version__)" diff --git a/tests/test_block_io.py b/tests/test_block_io.py index 0c27e85..8ba42d9 100644 --- a/tests/test_block_io.py +++ b/tests/test_block_io.py @@ -82,6 +82,7 @@ def test_mmap_reserve_copy_in_at_offset(self): def test_mmap_reserve_zero_length_raises(self): import persisting._core as _core + with pytest.raises(ValueError, match="length must be > 0"): _core.MmapRegion(0) diff --git a/tests/test_block_store.py b/tests/test_block_store.py index cb1c2bf..20c80a1 100644 --- a/tests/test_block_store.py +++ b/tests/test_block_store.py @@ -128,6 +128,7 @@ def test_partition_key_length_must_match_prefix_dims(self): class TestBlockStore: def test_get_triggers_fetch_from_l3(self): import numpy as np + l3 = NumpyBacking(SHAPE, dtype=np.float32) idx = (0, 0, 0, slice(0, 64)) l3.write(idx, np.ones(64, dtype=np.float32) * 42) @@ -147,6 +148,7 @@ def test_get_triggers_fetch_from_l3(self): def test_put_then_get(self): import numpy as np + store = BlockStore( DIMS, SHAPE, @@ -164,6 +166,7 @@ def test_put_then_get(self): def test_prefetch_and_wait(self): import numpy as np + l3 = NumpyBacking(SHAPE, dtype=np.float32) l3.write((0, 0, 0, slice(0, 128)), np.ones(128, dtype=np.float32)) store = BlockStore( @@ -185,6 +188,7 @@ def test_put_writes_through_to_l3(self): """put 写回 L3 后，另一 BlockStore 共享同一 L3 时能读到相同数据。""" import numpy as np import tempfile + with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as f: path = f.name try: @@ -213,6 +217,7 @@ def test_put_writes_through_to_l3(self): np.testing.assert_array_almost_equal(out, data) finally: import os + try: os.unlink(path) except OSError: @@ -228,6 +233,7 @@ def test_block_table_not_implemented_raises(self): def test_degraded_read_write_same_as_numpy(self): import numpy as np + backing = BlockMappedBacking((2, 2, 64), dtype=np.float32, block_tokens=64) idx = (0, 0, slice(0, 64)) data = np.ones(64, dtype=np.float32) * 7 @@ -239,9 +245,8 @@ def test_degraded_read_write_same_as_numpy(self): def test_use_mmap_read_write_same_as_numpy(self): """use_mmap=True 时 read/write 与降级实现结果一致。""" import numpy as np - backing = BlockMappedBacking( - (2, 2, 64), dtype=np.float32, block_tokens=64, use_mmap=True - ) + + backing = BlockMappedBacking((2, 2, 64), dtype=np.float32, block_tokens=64, use_mmap=True) idx = (0, 0, slice(0, 64)) data = np.ones(64, dtype=np.float32) * 7 backing.write(idx, data) @@ -254,6 +259,7 @@ class TestBlockStoreWithBlockMappedBacking: def test_get_triggers_fetch_from_l3(self): import numpy as np + l3 = NumpyBacking(SHAPE, dtype=np.float32) l3.write((0, 0, 0, slice(0, 64)), np.ones(64, dtype=np.float32) * 42) l1 = BlockMappedBacking(SHAPE, dtype=np.float32, block_tokens=BLOCK_TOKENS) @@ -274,6 +280,7 @@ def test_get_triggers_fetch_from_l3(self): def test_put_then_get(self): import numpy as np + l1 = BlockMappedBacking(SHAPE, dtype=np.float32, block_tokens=BLOCK_TOKENS) store = BlockStore( DIMS, @@ -295,6 +302,7 @@ def test_put_writes_through_to_l3(self): import numpy as np import os import tempfile + with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as f: path = f.name try: @@ -333,9 +341,8 @@ def test_put_writes_through_to_l3(self): def test_store_with_mmap_l1_put_then_get(self): """L1=BlockMappedBacking(use_mmap=True) 时 put 再 get 与 NumpyBacking 一致。""" import numpy as np - l1 = BlockMappedBacking( - SHAPE, dtype=np.float32, block_tokens=BLOCK_TOKENS, use_mmap=True - ) + + l1 = BlockMappedBacking(SHAPE, dtype=np.float32, block_tokens=BLOCK_TOKENS, use_mmap=True) store = BlockStore( DIMS, SHAPE, @@ -359,6 +366,7 @@ class TestPrefetchAsync: def test_wait_waits_for_background_prefetch(self): """prefetch(region) 后 wait(region) 阻塞直到后台预取完成，再 get 得到 L3 数据。""" import numpy as np + l3 = NumpyBacking(SHAPE, dtype=np.float32) l3.write((0, 0, 0, slice(0, 128)), np.ones(128, dtype=np.float32) * 99) store = BlockStore( @@ -380,6 +388,7 @@ def test_wait_waits_for_background_prefetch(self): def test_get_without_prefetch_sync_fetches(self): """不调用 prefetch，直接 get(region) 仍能同步从 L3 拉取并返回正确数据。""" import numpy as np + l3 = NumpyBacking(SHAPE, dtype=np.float32) l3.write((0, 0, 0, slice(0, 64)), np.arange(64, dtype=np.float32)) store = BlockStore( @@ -426,6 +435,7 @@ class TestBlockStoreWithRemoteBacking: def test_get_triggers_fetch_from_remote_stub(self): """L3 为 RemoteBacking(stub) 时，get 从 stub 拉块到 L1。""" import numpy as np + stub = {} idx_key = (0, 0, 0, (0, 64)) stub[idx_key] = np.ones(64, dtype=np.float32) * 42 @@ -452,6 +462,7 @@ def test_get_triggers_fetch_from_remote_stub(self): def test_put_writes_through_to_remote_stub(self): """put 写回 RemoteBacking(stub)，另一 BlockStore 共享同一 stub 时能读到。""" import numpy as np + stub = {} l3 = RemoteBacking( SHAPE, @@ -488,6 +499,7 @@ def test_open_tiered_returns_namespace(self): import numpy as np import persisting from persisting.core import Dimension + S = Dimension("session", "int") L = Dimension("layer", "int") T = Dimension("time", "int") diff --git a/tests/test_block_store_actor.py b/tests/test_block_store_actor.py index e7fa0a8..d83f713 100644 --- a/tests/test_block_store_actor.py +++ b/tests/test_block_store_actor.py @@ -29,6 +29,7 @@ def _pulsing_available(): try: import pulsing.core + return True except ImportError: return False @@ -39,6 +40,7 @@ def _pulsing_available(): async def test_block_store_actor_put_then_get(): """spawn BlockStoreActor，put_region 后 get_region 返回相同数据。""" import pulsing.core + await pulsing.core.init() try: BlockStoreActor = get_block_store_actor_class() @@ -51,6 +53,7 @@ async def test_block_store_actor_put_then_get(): ) try: import numpy as np + view = TensorView(DIMS) r = view["s0", 0, 0, 0:64] region_ser = region_serialize(r, DIMS) @@ -70,6 +73,7 @@ async def test_block_store_actor_put_then_get(): async def test_actor_store_get_async_put_async(): """ActorStore(proxy, dims) 的 get_async/put_async 与直接调 actor 一致。""" import pulsing.core + await pulsing.core.init() try: BlockStoreActor = get_block_store_actor_class() @@ -84,6 +88,7 @@ async def test_actor_store_get_async_put_async(): view = TensorView(DIMS) r = view["s0", 0, 0, 0:64] import numpy as np + data = np.ones(64, dtype=np.float32) * 7 await store.put_async(r, data) out = await store.get_async(r) @@ -101,6 +106,7 @@ def test_region_serialize_deserialize_roundtrip(): r2 = region_deserialize(ser, DIMS) # 用 BlockStore 验证：同一 region 应得到同一 block 列表 from persisting.store import region_to_blocks + blocks = region_to_blocks(r, DIMS, TIME, BLOCK_TOKENS, SHAPE) blocks2 = region_to_blocks(r2, DIMS, TIME, BLOCK_TOKENS, SHAPE) assert blocks == blocks2 diff --git a/tests/test_local_tensor_store.py b/tests/test_local_tensor_store.py index e458052..b767b83 100644 --- a/tests/test_local_tensor_store.py +++ b/tests/test_local_tensor_store.py @@ -8,13 +8,13 @@ import persisting from persisting.core import Dimension, TensorView from persisting.store import ( - Handler, - LocalTensorStore, - MmapBacking, - SafetensorsBacking, - TensorNamespace, - region_to_index, -) + Handler, + LocalTensorStore, + MmapBacking, + SafetensorsBacking, + TensorNamespace, + region_to_index, + ) except ImportError as e: pytest.skip(f"persisting._core 或 store 不可用: {e}", allow_module_level=True) @@ -208,6 +208,7 @@ def test_safetensors_backing_import_error_without_lib(self): """未安装 safetensors 时构造应给出明确 ImportError。""" try: import safetensors # noqa: F401 + pytest.skip("safetensors 已安装，跳过") except ImportError: pass diff --git a/tests/test_metadata.py b/tests/test_metadata.py index 43cc3cc..6043a77 100644 --- a/tests/test_metadata.py +++ b/tests/test_metadata.py @@ -6,12 +6,16 @@ def test_batch_meta_basic_properties(): SampleMeta( partition_id="p0", global_index=0, - fields={"x": FieldMeta(name="x", dtype="float32", shape=(2,), production_status="ready")}, + fields={ + "x": FieldMeta(name="x", dtype="float32", shape=(2,), production_status="ready") + }, ), SampleMeta( partition_id="p0", global_index=1, - fields={"x": FieldMeta(name="x", dtype="float32", shape=(2,), production_status="ready")}, + fields={ + "x": FieldMeta(name="x", dtype="float32", shape=(2,), production_status="ready") + }, ), ] meta = BatchMeta(samples=samples, custom_meta={0: {"k": "v"}}) @@ -27,7 +31,9 @@ def test_batch_meta_roundtrip(): SampleMeta( partition_id="p1", global_index=3, - fields={"y": FieldMeta(name="y", dtype="int64", shape=(1,), production_status="ready")}, + fields={ + "y": FieldMeta(name="y", dtype="int64", shape=(1,), production_status="ready") + }, ) ] ) diff --git a/tests/test_queue_tensordict.py b/tests/test_queue_tensordict.py index abae35f..232a722 100644 --- a/tests/test_queue_tensordict.py +++ b/tests/test_queue_tensordict.py @@ -38,7 +38,9 @@ async def test_queue_put_get_tensordict(tmp_path): assert meta is not None assert meta.size == 3 - sampled_meta = await queue.get_meta(fields=["x", "y"], batch_size=2, task_name="t1", partition_id="p0") + sampled_meta = await queue.get_meta( + fields=["x", "y"], batch_size=2, task_name="t1", partition_id="p0" + ) assert sampled_meta.size == 2 data = await queue.get_data(sampled_meta, partition_id="p0") assert len(data) == 2 or data.batch_size[0] == 2 diff --git a/tests/test_taa_subscript_planning.py b/tests/test_taa_subscript_planning.py index 0afd034..ff111a3 100644 --- a/tests/test_taa_subscript_planning.py +++ b/tests/test_taa_subscript_planning.py @@ -41,6 +41,7 @@ # 轻量访问：tensor[{SESSION: xx}, :, 0:100] # --------------------------------------------------------------------------- + class TestTensorViewAccess: """tensor[..., :, 0:100] 式访问，得到 Region。""" @@ -89,6 +90,7 @@ def test_tensor_wrong_length_raises(self): # Planning：canonicalize / project_prefix / is_point_query / is_range_query # --------------------------------------------------------------------------- + class TestTaaPlanning: """用 tensor 下标得到 Region 后做 planning。"""