hurray_io/file/writer.rs
1use hurray_core::TensorDescriptor;
2use tokio::io::{AsyncWrite, AsyncWriteExt};
3
4use crate::file::sans_io::{flatten_composite, FileWriteOp, FileWriterCore};
5// `pub use`, not `use`: 0.1.0 published this at `file::writer::FileCompositeNode` too (see
6// reader.rs). `hurray_io::file::FileCompositeNode` is canonical.
7pub use crate::file::types::FileCompositeNode;
8use crate::file::types::{FileWriterOptions, KvValue};
9use crate::stream::{Unusable, WriteStep};
10use crate::{Error, Result};
11
12/// Writes a Hurray file in a single forward pass without seeks.
13///
14/// The file format is:
15/// ```text
16/// [ File header ] 64 bytes
17/// [ Tensor 0 ] descriptor + padding + data buffers
18/// [ Tensor 1 ]
19/// ...
20/// [ KV section ] optional, written by finish()
21/// [ Footer index ] written by finish()
22/// [ Trailer ] 40 bytes, written by finish()
23/// ```
24///
25/// This is a driver, not an encoder: [`FileWriterCore`] holds the container format, the
26/// offsets and the growing index, and this type only writes the bytes the core offers and
27/// reports how many the sink took (ADR-041 § 7).
28///
29/// `HAS_KV_METADATA` is not set in the header because KV content is supplied to
30/// [`finish`][FileWriter::finish], after the header has already gone out, and this writer
31/// never seeks back. Readers treat a non-zero trailer `kv_offset` as authoritative.
32///
33/// # Examples
34///
35/// ```no_run
36/// # #[tokio::main]
37/// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
38/// use hurray_core::{
39/// BufferHandle, DeviceTag, ElementType, LayoutDescriptor, Shape, SyncMode,
40/// TensorDescriptor, MIN_BUFFER_ALIGNMENT,
41/// };
42/// use hurray_io::file::{FileWriter, KvValue};
43///
44/// let handle = BufferHandle::new(
45/// 16, MIN_BUFFER_ALIGNMENT, DeviceTag::Cpu, SyncMode::ProducerSynced,
46/// )?;
47/// let descriptor = TensorDescriptor::new(
48/// 1, 0, ElementType::Float32, Shape::new(vec![2u64, 2])?, 0,
49/// LayoutDescriptor::RowMajor, vec![handle], None, None, None, None,
50/// )?;
51/// let data = [0u8; 16];
52///
53/// let file = tokio::fs::File::create("model.hrry").await?;
54/// let mut writer = FileWriter::new(file).await?;
55/// writer.write_tensor("weight", &descriptor, &[&data]).await?;
56/// writer
57/// .finish(vec![("producer".into(), KvValue::String("hurray".into()))])
58/// .await?;
59/// # Ok(())
60/// # }
61/// ```
62pub struct FileWriter<W> {
63 inner: W,
64 core: FileWriterCore,
65}
66
67impl<W: AsyncWrite + Unpin> FileWriter<W> {
68 /// Creates a writer with default options (4096-byte buffer alignment, unsorted index).
69 pub async fn new(inner: W) -> Result<Self> {
70 Self::with_options(inner, FileWriterOptions::default()).await
71 }
72
73 /// Creates a writer with custom options.
74 ///
75 /// # Errors
76 ///
77 /// - [`Error::InvalidHeader`] — `data_buffer_alignment` is not a power of two in
78 /// `[4096, 2097152]`
79 /// - [`Error::Io`] — the header could not be written
80 pub async fn with_options(mut inner: W, options: FileWriterOptions) -> Result<Self> {
81 let mut core = FileWriterCore::with_options(options)?;
82 drain(&mut core.header()?, &mut inner).await?;
83 Ok(Self { inner, core })
84 }
85
86 /// Encodes and writes one tensor.
87 ///
88 /// # Errors
89 ///
90 /// - [`Error::TensorNameEmpty`] / [`Error::TensorNameTooLong`] — invalid name
91 /// - [`Error::DuplicateTensorName`] — name already written to this file
92 /// - [`Error::MultiBufferLengthMismatch`] — wrong buffer count
93 /// - [`Error::BufferSizeMismatch`] — buffer length ≠ handle `byte_size`
94 /// - [`Error::Core`] — descriptor encoding failed
95 /// - [`Error::SinkClosed`] — the sink stopped accepting bytes mid-tensor
96 /// - [`Error::Io`] — underlying write error
97 pub async fn write_tensor(
98 &mut self,
99 name: &str,
100 desc: &TensorDescriptor,
101 buffers: &[&[u8]],
102 ) -> Result<()> {
103 let mut op = self.core.tensor(name, desc, buffers)?;
104 drain(&mut op, &mut self.inner).await
105 }
106
107 /// Writes a composite tensor: its head, then every member's descriptor and data,
108 /// contiguously and in order (ADR-027 § Binding).
109 ///
110 /// Every tensor — the head and each member — gets its own footer-index entry, so all
111 /// are individually addressable by name via
112 /// [`read_tensor`][crate::file::FileReader::read_tensor]. Membership is recoverable by
113 /// [`read_composite`][crate::file::FileReader::read_composite] from the head's
114 /// `member_count` plus file-offset adjacency (the members are the tensors written
115 /// immediately after the head). Nested composites are written in the same pass.
116 ///
117 /// The whole group is validated up front — reusing
118 /// [`CompositeValidator`][hurray_core::CompositeValidator] for member count and
119 /// per-rule constraints (partition exact-cover, overlay ordering) — before any tensor
120 /// is written.
121 ///
122 /// # Errors
123 ///
124 /// - [`Error::Core`] — the head is not a valid composite head, or validation failed
125 /// - the name/buffer errors of [`write_tensor`][FileWriter::write_tensor] for the head
126 /// or any member
127 /// - [`Error::Io`] — underlying write error
128 ///
129 /// # Examples
130 ///
131 /// ```no_run
132 /// # #[tokio::main]
133 /// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
134 /// use hurray_core::{
135 /// layout::{CompositeLayout, CompositionRule, LayoutDescriptor},
136 /// ElementType, Shape, TensorDescriptor,
137 /// };
138 /// use hurray_io::file::{FileCompositeNode, FileWriter};
139 ///
140 /// let head = TensorDescriptor::new(
141 /// 1, 0, ElementType::Float32, Shape::new(vec![8u64, 8]).unwrap(), 0,
142 /// LayoutDescriptor::Composite(CompositeLayout::new(CompositionRule::Partition, 2).unwrap()),
143 /// vec![], None, None, None, None,
144 /// )?;
145 /// # let members: Vec<FileCompositeNode> = vec![];
146 /// let file = tokio::fs::File::create("model.hrry").await?;
147 /// let mut writer = FileWriter::new(file).await?;
148 /// writer.write_composite("weight", &head, &members).await?;
149 /// writer.finish(vec![]).await?;
150 /// # Ok(())
151 /// # }
152 /// ```
153 pub async fn write_composite(
154 &mut self,
155 head_name: &str,
156 head: &TensorDescriptor,
157 members: &[FileCompositeNode<'_>],
158 ) -> Result<()> {
159 // The tree is validated and flattened into write order first, so a torn or invalid
160 // composite never reaches the sink — and the nesting is walked with an explicit
161 // stack in the core rather than by recursing here (§ 10).
162 for (index, node) in flatten_composite(head_name, head, members)?
163 .into_iter()
164 .enumerate()
165 {
166 // The op is bound to a local and dropped before the poison below: it borrows
167 // the core mutably for its whole life, so `self.core` is unreachable until it
168 // is gone. That is also why `mark_unusable` on an op cannot serve here — a
169 // *submission* failure never produces one.
170 let result = match self.core.tensor(node.name, node.descriptor, node.buffers) {
171 Ok(mut op) => drain(&mut op, &mut self.inner).await,
172 Err(e) => Err(e),
173 };
174 if let Err(e) = result {
175 // A group that began and did not finish leaves a head with fewer than
176 // `member_count` members after it, which `file-format.md § Composite
177 // Tensors` calls invalid. No `finish` can rescue that, so the writer must
178 // not go on to write a clean index over it.
179 // `index > 0` is exactly "an earlier node completed": the loop returns on
180 // the first failure, so every node before this one is on the sink.
181 if index > 0 {
182 self.core.poison(Unusable::IncompleteOutput);
183 }
184 return Err(e);
185 }
186 }
187 Ok(())
188 }
189
190 /// Writes the KV section, footer index, and trailer, then flushes.
191 ///
192 /// `kv` is a list of `(key, value)` pairs. Keys must be non-empty, at most
193 /// 65 535 bytes, and unique within the list.
194 ///
195 /// # Errors
196 ///
197 /// - [`Error::KvKeyEmpty`] / [`Error::KvKeyTooLong`] / [`Error::DuplicateKvKey`] —
198 /// invalid KV keys
199 /// - [`Error::InvalidHeader`] — a KV array is empty, nested, or not homogeneous
200 /// - [`Error::SinkClosed`] — the sink stopped accepting bytes mid-footer
201 /// - [`Error::Io`] — underlying write error
202 pub async fn finish(mut self, kv: Vec<(String, KvValue)>) -> Result<W> {
203 let mut op = self.core.finish(kv)?;
204 drain(&mut op, &mut self.inner).await?;
205 self.inner.flush().await?;
206 Ok(self.inner)
207 }
208}
209
210/// Writes every byte of one submitted item to `sink`.
211///
212/// A free function rather than a method: the operation holds a mutable borrow of the
213/// writer's core for as long as it lives, so only disjoint *field* borrows let the sink be
214/// used at the same time.
215async fn drain<W: AsyncWrite + Unpin>(op: &mut FileWriteOp<'_>, sink: &mut W) -> Result<()> {
216 while let WriteStep::Output { .. } = op.step()? {
217 let n = {
218 let src = op.output_src();
219 loop {
220 // `write`, not `write_all`: a failed bulk write cannot report how many
221 // bytes it moved, and that count is the file position (§ 6).
222 match sink.write(src).await {
223 Ok(n) => break n,
224 // `write_all`, which this replaces, absorbed this internally.
225 Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
226 // This call put no bytes on the sink, so its count is known to be
227 // zero. That is not the same as the file being intact: if an earlier
228 // call in this item moved bytes, the core is already torn and the next
229 // call on this writer reports `CoreUnusable` (§ 6b).
230 Err(e) => return Err(Error::Io(e)),
231 }
232 }
233 };
234
235 if n == 0 {
236 // § 4: a zero-length write is never progress. A closed sink is its own signal,
237 // not an I/O error to retry.
238 return Err(op.report_sink_closed());
239 }
240 op.report_written(n)?;
241 }
242 Ok(())
243}