File size: 9,223 Bytes
a468182 2ca20a1 a468182 2ca20a1 a468182 2ca20a1 a468182 2ca20a1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 | """Resume validation and transactional publication of ordered embedding windows."""
from __future__ import annotations
from collections.abc import Sequence
from pathlib import Path
from typing import Any
from .identity import _RUN_FINGERPRINT_SCHEMA_VERSION
from .pooling import Pooler
from .storage import (
SafetensorsStreamWriter,
append_sqlite_records,
initialize_sqlite_run,
load_result,
load_sqlite_result,
safetensors_result_exists,
save_result,
tensor_sha256,
update_sqlite_run_metadata,
)
from .types import EmbeddingInput, EmbeddingRecord, EmbeddingResult, LazyTensorReference
def _output_exists(path: str | Path, format: str) -> bool:
path = Path(path)
if format == "sqlite":
return path.is_file()
return safetensors_result_exists(path)
def _output_descriptor(position: int, record: EmbeddingRecord) -> dict[str, Any]:
tensor = record.tensor
if isinstance(tensor, LazyTensorReference):
dtype = tensor.dtype
shape = tensor.shape
digest = tensor.sha256
else:
dtype = str(tensor.dtype).removeprefix("torch.")
shape = tuple(tensor.shape)
digest = tensor_sha256(tensor)
return {
"position": position,
"id": record.id,
"dtype": dtype,
"shape": shape,
"sha256": digest,
}
class EmbeddingOutput:
"""Own the resumable prefix and the commit state of one output destination."""
def __init__(
self,
records: Sequence[EmbeddingInput],
*,
output: str | Path | None,
format: str,
resume: bool,
shard_size: int,
run_fingerprint: str,
input_fingerprint: str,
model_state_fingerprint: str | None,
model_state_fingerprint_source: str,
pooler: Pooler | None,
pooling_names: Sequence[str],
) -> None:
self.output = output
self.format = format
self.shard_size = shard_size
self.completed: EmbeddingResult | None = None
output_already_exists = output is not None and _output_exists(output, format)
existing: EmbeddingResult | None = None
self.start_position = 0
if output is not None and resume and output_already_exists:
if format == "sqlite":
try:
existing = load_sqlite_result(output, run_id=run_fingerprint)
except KeyError:
existing = load_result(output, format=format)
else:
existing = load_result(output, format=format)
if existing.metadata.get("fingerprint_schema_version") != (
_RUN_FINGERPRINT_SCHEMA_VERSION
):
raise ValueError(
"Existing embeddings use an incompatible run fingerprint schema; "
"choose another output or set resume=False."
)
if existing.metadata.get("run_fingerprint") != run_fingerprint:
raise ValueError(
"Existing embeddings were produced by a different run fingerprint; "
"choose another output or set resume=False."
)
if len(existing) > len(records):
raise ValueError(
"Existing embeddings are not an ordered prefix of the requested inputs."
)
prefix_matches = all(
(observed.id, observed.sequence) == (expected.id, expected.sequence)
for expected, observed in zip(records, existing, strict=False)
)
if not prefix_matches:
raise ValueError(
"Existing embeddings are not an ordered prefix of the requested inputs."
)
if len(existing) == len(records) and existing.metadata.get("complete", True):
self.completed = existing
return
self.start_position = len(existing)
self.sqlite_run_id: str | None = None
self.sqlite_replace_on_first_commit = False
self.sqlite_initial_metadata: dict[str, Any] | None = None
if output is not None and format == "sqlite":
self.sqlite_initial_metadata = {
"format_version": 1,
"fingerprint_schema_version": _RUN_FINGERPRINT_SCHEMA_VERSION,
"run_fingerprint": run_fingerprint,
"input_fingerprint": input_fingerprint,
"model_state_fingerprint": model_state_fingerprint,
"model_state_fingerprint_source": model_state_fingerprint_source,
"complete": False,
}
self.sqlite_run_id = run_fingerprint
if not resume and output_already_exists:
prior_run_readable = True
try:
load_sqlite_result(output, run_id=run_fingerprint)
except KeyError:
prior_run_readable = False
if prior_run_readable:
# Keep an exact prior run readable until replacement inference
# has produced the first complete commit window.
self.sqlite_replace_on_first_commit = True
if not self.sqlite_replace_on_first_commit:
initialize_sqlite_run(
output,
self.sqlite_initial_metadata,
resume=resume,
)
stream_safetensors = output is not None and format == "safetensors"
self.output_records: list[EmbeddingRecord] = (
[] if self.sqlite_run_id is not None or stream_safetensors else list(existing or ())
)
self.output_descriptors: list[dict[str, Any]] | None = [] if output is None else None
self.pool_slices: dict[str, tuple[int, int]] = {}
if existing and pooler is not None:
pooled_width = existing[0].load_tensor().shape[-1]
if pooled_width % len(pooling_names) != 0:
raise ValueError("Stored pooled width is inconsistent with pooling metadata.")
self.pool_slices = pooler.output_slices(pooled_width // len(pooling_names))
self.safetensors_writer: SafetensorsStreamWriter | None = None
if stream_safetensors:
if output is None:
raise RuntimeError(
"Safetensors streaming was enabled without an output destination."
)
transactional_overwrite = output_already_exists and not resume
self.safetensors_writer = SafetensorsStreamWriter(
output,
{
"format_version": 1,
"fingerprint_schema_version": _RUN_FINGERPRINT_SCHEMA_VERSION,
"run_fingerprint": run_fingerprint,
"input_fingerprint": input_fingerprint,
"model_state_fingerprint": model_state_fingerprint,
"model_state_fingerprint_source": model_state_fingerprint_source,
"complete": False,
},
shard_size=shard_size,
existing=existing or (),
reuse_existing=bool(resume and existing is not None),
publish_initial=not transactional_overwrite,
publish_incremental=not transactional_overwrite,
)
def append(self, window_start: int, new_records: list[EmbeddingRecord]) -> None:
"""Commit a complete ordered window at the storage format's granularity."""
if self.output_descriptors is not None:
self.output_descriptors.extend(
_output_descriptor(window_start + offset, record)
for offset, record in enumerate(new_records)
)
if self.output is not None and self.sqlite_run_id is not None:
append_sqlite_records(
self.output,
self.sqlite_run_id,
window_start,
new_records,
replace_metadata=(
self.sqlite_initial_metadata if self.sqlite_replace_on_first_commit else None
),
)
self.sqlite_replace_on_first_commit = False
elif self.safetensors_writer is not None:
self.safetensors_writer.append(new_records)
else:
self.output_records.extend(new_records)
def finish(self, metadata: dict[str, Any]) -> EmbeddingResult:
"""Publish completion only after every window has committed."""
if self.output is not None and self.sqlite_run_id is not None:
update_sqlite_run_metadata(self.output, self.sqlite_run_id, metadata)
return load_sqlite_result(self.output, run_id=self.sqlite_run_id)
if self.safetensors_writer is not None:
return self.safetensors_writer.publish(complete=True, metadata=metadata)
embedding_result = EmbeddingResult(self.output_records, metadata)
if self.output is not None:
return save_result(
embedding_result, self.output, format=self.format, shard_size=self.shard_size
)
return embedding_result
|