diff --git a/.gitignore b/.gitignore index cf39a7699..99fb49f35 100644 --- a/.gitignore +++ b/.gitignore @@ -22,3 +22,4 @@ venv **/*.env .worktrees/ **/node_modules/ +**/__pycache__/ diff --git a/Cargo.lock b/Cargo.lock index 7e281275f..ec11d6d27 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1921,7 +1921,10 @@ dependencies = [ "async-std", "async-trait", "bytes", + "chrono", "clap", + "ctor", + "dirs", "futures", "http", "more-asserts", @@ -1929,6 +1932,7 @@ dependencies = [ "serde", "serde_json", "serial_test", + "sha2 0.11.0", "smol", "tempfile", "thiserror 2.0.18", diff --git a/hf_xet/src/lib.rs b/hf_xet/src/lib.rs index e51bd46f7..4b5a0520e 100644 --- a/hf_xet/src/lib.rs +++ b/hf_xet/src/lib.rs @@ -8,6 +8,8 @@ mod py_download_stream_handle; mod py_file_download_group; mod py_file_download_handle; mod py_file_upload_handle; +mod py_range_upload_commit; +mod py_range_upload_edit; mod py_stream_upload_handle; mod py_upload_commit; mod py_xet_session; @@ -71,6 +73,8 @@ pub fn hf_xet(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/hf_xet/src/py_range_upload_commit.rs b/hf_xet/src/py_range_upload_commit.rs new file mode 100644 index 000000000..18c9399cb --- /dev/null +++ b/hf_xet/src/py_range_upload_commit.rs @@ -0,0 +1,262 @@ +//! Python bindings for XetRangeUploadCommit. + +use std::collections::HashMap; +use std::sync::{Arc, RwLock}; + +use pyo3::prelude::*; +use xet_pkg::xet_session::{ + GroupProgressReport, ItemProgressReport, XetRangeUploadCommit, XetRangeUploadEdit, XetRangeUploadReport, + XetTaskState, +}; +use xet_runtime::utils::UniqueId; + +use super::py_range_upload_edit::PyXetRangeUploadEdit; +use crate::background_progress::BackgroundProgress; +use crate::headers::build_header_map; +use crate::utils::{blocking_call_with_signal_check, convert_xet_error, task_state_display, task_state_to_pystate}; + +// ── Helpers ────────────────────────────────────────────────────────────────── + +fn item_reports_from_edit_handles( + handles: &Arc>>, +) -> HashMap { + // XetRangeUploadEdit doesn't expose progress directly like XetFileUpload does. + // For now, return an empty map. + handles + .read() + .map(|g| g.iter().filter_map(|_| None).collect()) + .unwrap_or_default() +} + +// ── build_range_upload_commit ──────────────────────────────────────────────── + +pub(crate) fn build_range_upload_commit( + py: Python<'_>, + session: &xet_pkg::xet_session::XetSession, + original_hash: String, + original_size: u64, + endpoint: Option, + token: Option, + token_expiry_unix_secs: Option, + token_refresh_url: Option, + token_refresh_headers: Option>, + custom_headers: Option>, + progress_callback: Option>, + progress_interval_ms: u64, +) -> PyResult { + let mut builder = session.new_range_upload().map_err(convert_xet_error)?; + if let Some(ep) = endpoint { + builder = builder.with_endpoint(&ep); + } + if let (Some(tok), Some(exp)) = (token, token_expiry_unix_secs) { + builder = builder.with_token_info(tok, exp); + } + if let Some(url) = token_refresh_url { + let headers = build_header_map(token_refresh_headers.unwrap_or_default())?; + builder = builder.with_token_refresh_url(url, headers); + } + + // custom_headers are already merged by PyXetSession via with_custom_headers + if let Some(headers) = custom_headers { + let hm = build_header_map(headers)?; + builder = builder.with_custom_headers(hm); + } + + let commit = py.detach(move || builder.build_blocking(original_hash, original_size).map_err(convert_xet_error))?; + + let (edit_handles, progress) = if let Some(callback) = progress_callback { + let handles: Arc>> = Arc::new(RwLock::new(Vec::new())); + let inner = commit.clone(); + let handles_for_thread = handles.clone(); + let progress = BackgroundProgress::spawn(py, callback, progress_interval_ms, move || { + let is_terminal = !matches!(inner.status(), Ok(XetTaskState::Running) | Ok(XetTaskState::Finalizing)); + let item_reports = inner.item_reports_from_upload_session(); + (inner.progress(), item_reports, is_terminal) + }); + (Some(handles), Some(progress)) + } else { + (None, None) + }; + + Ok(PyXetRangeUploadCommit { + inner: commit, + edit_handles, + progress, + }) +} + +// ── PyXetRangeUploadCommit ─────────────────────────────────────────────────── + +/// A commit that edits an existing file by uploading only changed byte ranges. +/// +/// Implements the context-manager protocol. Entering the ``with`` block returns +/// the commit itself. On normal exit :meth:`commit` is called automatically; +/// on exception :meth:`abort` is called automatically. +/// +/// ```python +/// with session.new_range_upload(hash, size) as commit: +/// commit.edit(1000, 2000).write(b"new data") +/// commit.append(500) +/// # commit() is called automatically on normal exit. +/// ``` +#[pyclass(name = "XetRangeUploadCommit")] +pub struct PyXetRangeUploadCommit { + pub(crate) inner: XetRangeUploadCommit, + /// Per-edit handles shared with the progress thread; None when no callback was registered. + edit_handles: Option>>>, + /// Background thread that polls progress and invokes the Python callback. + progress: Option, +} + +#[pymethods] +impl PyXetRangeUploadCommit { + fn __repr__(&self) -> String { + let status = task_state_display(self.inner.status()); + format!("XetRangeUploadCommit(status=\"{}\")", status) + } + + // ── Context manager ────────────────────────────────────────────────── + + fn __enter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> { + slf + } + + fn __exit__( + &self, + py: Python<'_>, + exc_type: Bound<'_, pyo3::PyAny>, + _exc_val: Bound<'_, pyo3::PyAny>, + _exc_tb: Bound<'_, pyo3::PyAny>, + ) -> PyResult { + if exc_type.is_none() { + // Normal exit: commit (signal-interruptible). + self.commit(py)?; + } else { + if let Err(e) = self.abort(py) { + tracing::warn!("abort() failed during __exit__ exception path: {e}"); + } + } + Ok(false) + } + + // ── Edit methods ───────────────────────────────────────────────────── + + /// Start a new edit: replace ``original_range`` with ``new_length`` bytes. + /// + /// Returns an :class:`XetRangeUploadEdit` handle. Feed data incrementally + /// with :meth:`XetRangeUploadEdit.write`. Call :meth:`XetRangeUploadEdit.finish` + /// before committing, or call :meth:`commit` directly — it will finalise all + /// pending edits automatically. + /// + /// The ``original_range`` parameter accepts a tuple ``(start, end)``. + pub fn edit(&self, original_range: (u64, u64), new_length: u64) -> PyResult { + let inner = self.inner.clone(); + let edit = inner.edit(original_range.0..original_range.1, new_length); + if let Some(ref handles) = self.edit_handles { + handles + .write() + .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))? + .push(edit.clone()); + } + Ok(PyXetRangeUploadEdit { inner: edit }) + } + + /// Convenience: insert ``new_length`` bytes at position ``pos``. + /// + /// Equivalent to ``edit((pos, pos), new_length)``. + pub fn insert(&self, pos: u64, new_length: u64) -> PyResult { + let inner = self.inner.clone(); + let edit = inner.insert(pos, new_length); + if let Some(ref handles) = self.edit_handles { + handles + .write() + .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))? + .push(edit.clone()); + } + Ok(PyXetRangeUploadEdit { inner: edit }) + } + + /// Convenience: delete bytes at ``start..end``. + /// + /// Equivalent to ``edit((start, end), 0)``. + pub fn delete(&self, start: u64, end: u64) -> PyResult { + let inner = self.inner.clone(); + let edit = inner.delete(start, end); + if let Some(ref handles) = self.edit_handles { + handles + .write() + .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))? + .push(edit.clone()); + } + Ok(PyXetRangeUploadEdit { inner: edit }) + } + + /// Convenience: append ``new_length`` bytes at the end of the file. + /// + /// Equivalent to ``edit((original_size, original_size), new_length)``. + pub fn append(&self, new_length: u64) -> PyResult { + let inner = self.inner.clone(); + let edit = inner.append(new_length); + if let Some(ref handles) = self.edit_handles { + handles + .write() + .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))? + .push(edit.clone()); + } + Ok(PyXetRangeUploadEdit { inner: edit }) + } + + // ── Commit / abort ─────────────────────────────────────────────────── + + /// Wait for all edits to be committed and return the result. + /// + /// Returns an :class:`XetRangeUploadReport` containing the composed file's + /// :class:`XetFileInfo`. Also called automatically when exiting a ``with`` + /// block without an exception. + /// + /// Releases the GIL while waiting, polling for ``KeyboardInterrupt`` every + /// 100 ms so that Ctrl-C is delivered promptly. + pub fn commit(&self, py: Python<'_>) -> PyResult { + let inner = self.inner.clone(); + let result = blocking_call_with_signal_check(py, move || inner.commit_blocking()); + if let (Some(handles), Some(progress)) = (&self.edit_handles, &self.progress) { + let progress_join_ret = if result.is_ok() { + progress.stop_and_emit(py, || { + let _item_reports = item_reports_from_edit_handles(handles); + (self.inner.progress(), _item_reports) + }) + } else { + progress.stop_and_join(py) + }; + if let Err(e) = progress_join_ret { + tracing::warn!(error = ?e, "PyXetRangeUploadCommit progress thread join failed"); + } + } + result + } + + /// Cancel all pending edits. + pub fn abort(&self, py: Python<'_>) -> PyResult<()> { + if let Some(progress) = &self.progress { + let _ = progress.stop_and_join(py); + } + self.inner.abort().map_err(convert_xet_error) + } + + // ── Progress / status ──────────────────────────────────────────────── + + /// Aggregate progress for this commit. + pub fn progress(&self) -> GroupProgressReport { + self.inner.progress() + } + + /// Get item reports from the upload session. + pub fn item_reports_from_upload_session(&self) -> std::collections::HashMap { + self.inner.item_reports_from_upload_session() + } + + /// Current task state as a :class:`XetTaskState` enum value. Raises on error. + pub fn status(&self) -> PyResult { + task_state_to_pystate(self.inner.status()) + } +} diff --git a/hf_xet/src/py_range_upload_edit.rs b/hf_xet/src/py_range_upload_edit.rs new file mode 100644 index 000000000..cc068748f --- /dev/null +++ b/hf_xet/src/py_range_upload_edit.rs @@ -0,0 +1,41 @@ +//! Python bindings for XetRangeUploadEdit. + +use std::sync::Arc; + +use pyo3::prelude::*; +use xet_pkg::xet_session::XetRangeUploadEdit; + +/// Handle for a single edit within a :class:`XetRangeUploadCommit`. +/// +/// Returned by :meth:`XetRangeUploadCommit.edit`, :meth:`XetRangeUploadCommit.insert`, +/// :meth:`XetRangeUploadCommit.delete`, and :meth:`XetRangeUploadCommit.append`. +/// Feed data incrementally with :meth:`write`. Call :meth:`finish` to finalise +/// the edit before calling :meth:`XetRangeUploadCommit.commit`, or call +/// :meth:`commit` directly — it will finalise all pending edits automatically. +#[pyclass(name = "XetRangeUploadEdit")] +pub struct PyXetRangeUploadEdit { + pub(crate) inner: XetRangeUploadEdit, +} + +#[pymethods] +impl PyXetRangeUploadEdit { + /// Feed data into this edit. + /// + /// May be called any number of times before :meth:`finish` or :meth:`XetRangeUploadCommit.commit`. + pub fn write(&self, data: Vec) { + self.inner.write(&data); + } + + /// Finalise the edit. + /// + /// This is optional: calling :meth:`XetRangeUploadCommit.commit` will finalise all + /// pending edits automatically. After a successful finish, subsequent calls return + /// ``None``. + pub fn finish(&self) -> PyResult<()> { + let inner = Arc::new(self.inner.clone()); + match inner.finish() { + Ok(_) => Ok(()), + Err(_) => Err(pyo3::exceptions::PyRuntimeError::new_err("edit was already finished")), + } + } +} diff --git a/hf_xet/src/py_xet_session.rs b/hf_xet/src/py_xet_session.rs index 24cab3668..7e128e035 100644 --- a/hf_xet/src/py_xet_session.rs +++ b/hf_xet/src/py_xet_session.rs @@ -6,6 +6,7 @@ use xet_runtime::config::XetConfig; use crate::py_download_stream_group::{PyXetDownloadStreamGroup, build_download_stream_group}; use crate::py_file_download_group::{PyXetFileDownloadGroup, build_file_download_group}; +use crate::py_range_upload_commit::{PyXetRangeUploadCommit, build_range_upload_commit}; use crate::py_upload_commit::{PyXetUploadCommit, build_upload_commit}; use crate::utils::{task_state_display, task_state_to_pystate}; use crate::{PyXetTaskState, convert_xet_error}; @@ -261,6 +262,71 @@ impl PyXetSession { ) } + /// Create a :class:`XetRangeUploadCommit` for editing an existing file by uploading + /// only changed byte ranges. + /// + /// A range upload (also called a "dirty upload") edits an existing file by uploading + /// only the changed byte ranges. The untouched regions are pulled from the original + /// file in CAS. + /// + /// Configure the commit with any combination of: + /// - ``endpoint`` — CAS server URL + /// - ``token`` — CAS token + /// - ``token_expiry_unix_secs`` — Token expiry as Unix timestamp (seconds) + /// - ``token_refresh_url`` — URL to refresh the token + /// - ``token_refresh_headers`` — Headers for the token refresh request + /// - ``custom_headers`` — Extra HTTP headers forwarded with every CAS request + /// - ``progress_callback`` — Python callable invoked periodically with progress info + /// - ``progress_interval_ms`` — Callback interval in milliseconds + /// + /// Then call :meth:`XetRangeUploadCommit.edit`, :meth:`XetRangeUploadCommit.insert`, + /// :meth:`XetRangeUploadCommit.delete`, or :meth:`XetRangeUploadCommit.append` to queue + /// edits. Feed data incrementally with :meth:`XetRangeUploadEdit.write`. Call + /// :meth:`XetRangeUploadEdit.finish` to finalise the edit before committing, or + /// call :meth:`XetRangeUploadCommit.commit` directly — it will finalise all + /// pending edits automatically. + #[pyo3(signature = ( + original_hash, + original_size, + endpoint = None, + token = None, + token_expiry_unix_secs = None, + token_refresh_url = None, + token_refresh_headers = None, + custom_headers = None, + progress_callback = None, + progress_interval_ms = 100, + ))] + pub fn new_range_upload( + &self, + py: Python<'_>, + original_hash: String, + original_size: u64, + endpoint: Option, + token: Option, + token_expiry_unix_secs: Option, + token_refresh_url: Option, + token_refresh_headers: Option>, + custom_headers: Option>, + progress_callback: Option>, + progress_interval_ms: u64, + ) -> PyResult { + build_range_upload_commit( + py, + &self.inner, + original_hash, + original_size, + endpoint, + token, + token_expiry_unix_secs, + token_refresh_url, + token_refresh_headers, + custom_headers, + progress_callback, + progress_interval_ms, + ) + } + /// Current task state as a :class:`XetTaskState` enum value. Raises on error. pub fn status(&self) -> PyResult { task_state_to_pystate(self.inner.status()) diff --git a/hf_xet/tests/test_range_upload.py b/hf_xet/tests/test_range_upload.py new file mode 100644 index 000000000..09e73b229 --- /dev/null +++ b/hf_xet/tests/test_range_upload.py @@ -0,0 +1,230 @@ +""" +End-to-end tests for XetRangeUploadCommit using a local CAS endpoint. + +Run after building the extension: + cd hf_xet && maturin develop + pytest tests/test_range_upload.py -v --tb=short + +These tests use a local CAS endpoint (no HF token or bucket required). +""" + +import hf_xet + + +# ── Helpers ────────────────────────────────────────────────────────────────── + +def _upload_bytes(session: hf_xet.XetSession, endpoint: str, data: bytes) -> hf_xet.XetFileInfo: + """Upload raw bytes via the regular upload commit and return XetFileInfo.""" + commit = session.new_upload_commit(endpoint=endpoint) + h = commit.start_upload_bytes(data, sha256=hf_xet.SKIP_SHA256) + commit.wait_to_finish() + return h.result().xet_info + + +def _download_via_group(session: hf_xet.XetSession, endpoint: str, file_info: hf_xet.XetFileInfo, dest_path: str): + """Download a file via the file download group and return the file info.""" + group = session.new_file_download_group(endpoint=endpoint) + h = group.start_download_file(file_info, dest_path) + report = group.wait_to_finish() + return report.downloads[h.task_id()].file_info + + +# ── Edit (replace bytes) ──────────────────────────────────────────────────── + +class TestRangeUploadEdit: + """Test: upload original, edit bytes 0..13, verify composed result.""" + + def test_e2e_range_upload_edit(self, endpoint, tmp_path): + original_data = b"Hello, World! This is a test file for range upload." + assert len(original_data) == 51 + + session = hf_xet.XetSession() + + # Step 1: upload original + original_info = _upload_bytes(session, endpoint, original_data) + assert original_info.file_size == 51 + + # Step 2: edit bytes 0..13 ("Hello, World!") -> "Universe! " (10 bytes) + # Expected new size: 51 - 13 + 10 = 48 + commit = session.new_range_upload(original_info.hash, original_info.file_size, endpoint=endpoint) + + edit = commit.edit((0, 13), 10) + edit.write(b"Universe! ") + # No need to call finish() explicitly — commit handles it + + report = commit.commit() + assert report.file_info.file_size == 48 + + # Step 3: download and verify content + dest_path = tmp_path / "range_upload_edit_test.txt" + _download_via_group(session, endpoint, report.file_info, str(dest_path)) + content = dest_path.read_bytes() + assert content == b"Universe! This is a test file for range upload." + + +# ── Insert ────────────────────────────────────────────────────────────────── + +class TestRangeUploadInsert: + """Test: upload original, insert bytes at position, verify result.""" + + def test_e2e_range_upload_insert(self, endpoint, tmp_path): + original_data = b"ABCDEF" + assert len(original_data) == 6 + + session = hf_xet.XetSession() + + # Step 1: upload original + original_info = _upload_bytes(session, endpoint, original_data) + assert original_info.file_size == 6 + + # Step 2: insert "XYZ" at position 2 (between B and C) + # Expected new size: 6 + 3 = 9 + commit = session.new_range_upload(original_info.hash, original_info.file_size, endpoint=endpoint) + + edit = commit.insert(2, 3) + edit.write(b"XYZ") + # No need to call finish() explicitly — commit handles it + + report = commit.commit() + assert report.file_info.file_size == 9 + + # Step 3: download and verify content + dest_path = tmp_path / "range_upload_insert_test.txt" + _download_via_group(session, endpoint, report.file_info, str(dest_path)) + content = dest_path.read_bytes() + assert content == b"ABXYZCDEF" + + +# ── Delete ────────────────────────────────────────────────────────────────── + +class TestRangeUploadDelete: + """Test: upload original, delete bytes, verify result.""" + + def test_e2e_range_upload_delete(self, endpoint, tmp_path): + original_data = b"Hello, World!" + assert len(original_data) == 13 + + session = hf_xet.XetSession() + + # Step 1: upload original + original_info = _upload_bytes(session, endpoint, original_data) + assert original_info.file_size == 13 + + # Step 2: delete bytes 5..12 (", World") — 7 bytes removed + # Expected new size: 13 - 7 = 6 + commit = session.new_range_upload(original_info.hash, original_info.file_size, endpoint=endpoint) + + edit = commit.delete(5, 12) + # No need to call finish() — delete edits have no data to write + + report = commit.commit() + assert report.file_info.file_size == 6 + + # Step 3: download and verify content + dest_path = tmp_path / "range_upload_delete_test.txt" + _download_via_group(session, endpoint, report.file_info, str(dest_path)) + content = dest_path.read_bytes() + assert content == b"Hello!" + + +# ── Append ────────────────────────────────────────────────────────────────── + +class TestRangeUploadAppend: + """Test: upload original, append bytes at end, verify result.""" + + def test_e2e_range_upload_append(self, endpoint, tmp_path): + original_data = b"Hello, " + assert len(original_data) == 7 + + session = hf_xet.XetSession() + + # Step 1: upload original + original_info = _upload_bytes(session, endpoint, original_data) + assert original_info.file_size == 7 + + # Step 2: append "World!" (6 bytes) at end + # Expected new size: 7 + 6 = 13 + commit = session.new_range_upload(original_info.hash, original_info.file_size, endpoint=endpoint) + + edit = commit.append(6) + edit.write(b"World!") + # No need to call finish() explicitly — commit handles it + + report = commit.commit() + assert report.file_info.file_size == 13 + + # Step 3: download and verify content + dest_path = tmp_path / "range_upload_append_test.txt" + _download_via_group(session, endpoint, report.file_info, str(dest_path)) + content = dest_path.read_bytes() + assert content == b"Hello, World!" + + +# ── Multiple edits ────────────────────────────────────────────────────────── + +class TestRangeUploadMultipleEdits: + """Test: upload original, apply multiple edits in one commit.""" + + def test_e2e_range_upload_multiple_edits(self, endpoint, tmp_path): + original_data = b"0123456789ABCDEF" # 16 bytes + session = hf_xet.XetSession() + + # Step 1: upload original + original_info = _upload_bytes(session, endpoint, original_data) + assert original_info.file_size == 16 + + # Step 2: apply multiple edits + # - edit 0..4 ("0123" -> "XXXX") - keep same length + # - insert 8, 3 ("---") - 3 bytes added + # - delete 14..16 ("EF") — 2 bytes removed + # Expected: "XXXX4567---89ABCD" = 17 bytes + commit = session.new_range_upload(original_info.hash, original_info.file_size, endpoint=endpoint) + + edit1 = commit.edit((0, 4), 4) + edit1.write(b"XXXX") + + edit2 = commit.insert(8, 3) + edit2.write(b"---") + + edit3 = commit.delete(14, 16) + + report = commit.commit() + assert report.file_info.file_size == 17 + + # Step 3: download and verify + dest_path = tmp_path / "range_upload_multi_test.txt" + _download_via_group(session, endpoint, report.file_info, str(dest_path)) + content = dest_path.read_bytes() + assert content == b"XXXX4567---89ABCD" + + def test_e2e_range_upload_multiple_edits_unsorted(self, endpoint, tmp_path): + original_data = b"0123456789ABCDEF" # 16 bytes + session = hf_xet.XetSession() + + # Step 1: upload original + original_info = _upload_bytes(session, endpoint, original_data) + assert original_info.file_size == 16 + + # Step 2: apply multiple edits + # - delete 14..16 ("EF") — 2 bytes removed + # - insert 8, 3 ("---") - 3 bytes added + # - edit 0..4 ("0123" -> "XXXX") - keep same length + # Expected: "XXXX4567---89ABCD" = 17 bytes + commit = session.new_range_upload(original_info.hash, original_info.file_size, endpoint=endpoint) + + edit3 = commit.delete(14, 16) + + edit2 = commit.insert(8, 3) + edit2.write(b"---") + + edit1 = commit.edit((0, 4), 4) + edit1.write(b"XXXX") + + report = commit.commit() + assert report.file_info.file_size == 17 + + # Step 3: download and verify + dest_path = tmp_path / "range_upload_multi_test.txt" + _download_via_group(session, endpoint, report.file_info, str(dest_path)) + content = dest_path.read_bytes() + assert content == b"XXXX4567---89ABCD" diff --git a/xet_data/src/processing/range_upload.rs b/xet_data/src/processing/range_upload.rs index 042c74080..76d6228f2 100644 --- a/xet_data/src/processing/range_upload.rs +++ b/xet_data/src/processing/range_upload.rs @@ -90,6 +90,7 @@ pub async fn upload_ranges( original_hash: MerkleHash, original_size: u64, mut dirty_inputs: Vec, + upload_session: Option>, ) -> Result { validate_dirty_ranges(&dirty_inputs, original_size)?; let total_size = compute_total_size(original_size, &dirty_inputs)?; @@ -184,7 +185,10 @@ pub async fn upload_ranges( let gap_verification = response.gap_verification; let ctx = config.ctx.clone(); - let session = FileUploadSession::new(config.clone()).await?; + let session: Arc = match upload_session { + Some(upload_session) => upload_session, + None => FileUploadSession::new(config.clone()).await?, + }; let mut input_idx = 0usize; let mut uploaded: Vec = Vec::with_capacity(response.windows.len()); @@ -668,6 +672,7 @@ mod tests { original_hash, original_size, make_legacy_inputs(&[(dirty_start as u64, dirty_end as u64)], &modified_data, original_size, total_size), + None, ) .await .unwrap(); @@ -713,6 +718,7 @@ mod tests { original_hash, original_size, make_legacy_inputs(&[], &[], original_size, truncated_size), + None, ) .await .unwrap(); @@ -753,6 +759,7 @@ mod tests { original_hash, original_size, make_legacy_inputs(&[(original_size, total_size)], &full_data, original_size, total_size), + None, ) .await .unwrap(); @@ -792,6 +799,7 @@ mod tests { original_hash, original_size, make_legacy_inputs(&[(0, 4096)], &modified_data, original_size, total_size), + None, ) .await .unwrap(); @@ -832,6 +840,7 @@ mod tests { original_hash, original_size, make_legacy_inputs(&[(10_000, 12_000), (200_000, 202_000)], &modified_data, original_size, total_size), + None, ) .await .unwrap(); @@ -878,6 +887,7 @@ mod tests { original_hash, original_size, make_legacy_inputs(&[(original_size, total_size)], &full_data, original_size, total_size), + None, ) .await .unwrap(); @@ -920,6 +930,7 @@ mod tests { original_hash, original_size, make_legacy_inputs(&[(original_size, total_size)], &sparse_staging, original_size, total_size), + None, ) .await .unwrap(); @@ -1020,6 +1031,7 @@ mod tests { original_hash, size, make_legacy_inputs(&[(boundary, dirty_end)], &expected, size, size), + None, ) .await .unwrap(); @@ -1043,7 +1055,7 @@ mod tests { let data = random_data(70, 256 * 1024); let hash = upload_file(&config, &data).await; let size = data.len() as u64; - let result = upload_ranges(config, cas_client, hash, size, make_legacy_inputs(&[], &[], size, size)) + let result = upload_ranges(config, cas_client, hash, size, make_legacy_inputs(&[], &[], size, size), None) .await .unwrap(); @@ -1062,7 +1074,7 @@ mod tests { let data = random_data(71, 256 * 1024); let hash = upload_file(&config, &data).await; let size = data.len() as u64; - let err = upload_ranges(config, cas_client, hash, size, make_dummy_inputs(&[(100, size + 1)])).await; + let err = upload_ranges(config, cas_client, hash, size, make_dummy_inputs(&[(100, size + 1)]), None).await; assert!(err.is_err(), "dirty range past total_size should be rejected"); } @@ -1076,7 +1088,8 @@ mod tests { let data = random_data(60, 256 * 1024); let hash = upload_file(&config, &data).await; let size = data.len() as u64; - let err = upload_ranges(config, cas_client, hash, size, make_dummy_inputs(&[(100, 300), (200, 400)])).await; + let err = + upload_ranges(config, cas_client, hash, size, make_dummy_inputs(&[(100, 300), (200, 400)]), None).await; assert!(err.is_err(), "overlapping ranges should be rejected"); } @@ -1105,7 +1118,7 @@ mod tests { reader: Box::pin(Cursor::new(vec![0xBB; 10])), }, ]; - let err = upload_ranges(config, cas_client, original_hash, 0, inputs).await; + let err = upload_ranges(config, cas_client, original_hash, 0, inputs, None).await; assert!(err.is_err(), "ranges with end > original_size must be rejected for empty originals too"); } @@ -1119,7 +1132,8 @@ mod tests { let data = random_data(62, 256 * 1024); let hash = upload_file(&config, &data).await; let size = data.len() as u64; - let err = upload_ranges(config, cas_client, hash, size, make_dummy_inputs(&[(300, 400), (100, 200)])).await; + let err = + upload_ranges(config, cas_client, hash, size, make_dummy_inputs(&[(300, 400), (100, 200)]), None).await; assert!(err.is_err(), "unsorted ranges should be rejected"); } @@ -1151,6 +1165,7 @@ mod tests { original_hash, original_size, make_legacy_inputs(&[(dirty_start, dirty_end)], &modified, original_size, original_size), + None, ) .await .unwrap(); @@ -1187,9 +1202,10 @@ mod tests { reader: Box::pin(Cursor::new(dirty_data.to_vec())), }]; - let result = upload_ranges(config.clone(), cas_client.clone(), original_hash, original_size, dirty_inputs) - .await - .unwrap(); + let result = + upload_ranges(config.clone(), cas_client.clone(), original_hash, original_size, dirty_inputs, None) + .await + .unwrap(); assert_eq!(result.file_size(), Some(original_size)); @@ -1229,6 +1245,7 @@ mod tests { original_hash, original_size, make_legacy_inputs(&[], &[], original_size, truncated_size), + None, ) .await .unwrap(); @@ -1282,6 +1299,7 @@ mod tests { original_hash, original_size, make_legacy_inputs(&[(dirty_start, dirty_end)], &staging, original_size, truncated_size), + None, ) .await .unwrap(); @@ -1472,7 +1490,7 @@ mod tests { ) { let original_hash = upload_file(config, original).await; let original_size = original.len() as u64; - let result = upload_ranges(config.clone(), cas_client.clone(), original_hash, original_size, inputs) + let result = upload_ranges(config.clone(), cas_client.clone(), original_hash, original_size, inputs, None) .await .unwrap(); assert_eq!(result.file_size(), Some(expected.len() as u64), "file size mismatch"); @@ -1522,7 +1540,7 @@ mod tests { inputs.sort_by_key(|d| d.original_range.start); } - let result = upload_ranges(config.clone(), cas_client.clone(), original_hash, original_size, inputs) + let result = upload_ranges(config.clone(), cas_client.clone(), original_hash, original_size, inputs, None) .await .unwrap(); @@ -1569,7 +1587,7 @@ mod tests { reader: Box::pin(Cursor::new(append_extra)), }, ]; - let result = upload_ranges(config.clone(), cas_client.clone(), original_hash, original_size, inputs) + let result = upload_ranges(config.clone(), cas_client.clone(), original_hash, original_size, inputs, None) .await .unwrap(); @@ -1598,7 +1616,7 @@ mod tests { new_length: total_size, reader: Box::pin(Cursor::new(new_data.clone())), }]; - let result = upload_ranges(config.clone(), cas_client.clone(), original_hash, 0, inputs) + let result = upload_ranges(config.clone(), cas_client.clone(), original_hash, 0, inputs, None) .await .unwrap(); @@ -1627,6 +1645,7 @@ mod tests { original_hash, original_size, make_legacy_inputs(&[], &[], original_size, 0), + None, ) .await .unwrap(); @@ -1951,9 +1970,10 @@ mod tests { let edits = build_random_non_overlapping_edits(&mut rng, expected.len(), 8); let expected_next = apply_planned_edits(&expected, &edits); let inputs = edits_to_dirty_inputs(&edits); - let result = upload_ranges(config.clone(), cas_client.clone(), original_hash, original_size, inputs) - .await - .unwrap(); + let result = + upload_ranges(config.clone(), cas_client.clone(), original_hash, original_size, inputs, None) + .await + .unwrap(); let result_hash = MerkleHash::from_hex(result.hash()).unwrap(); assert_eq!( @@ -1993,7 +2013,7 @@ mod tests { let edits_summary = summarize_edits(&edits); let expected_next = apply_planned_edits(&expected, &edits); let inputs = edits_to_dirty_inputs(&edits); - let result = upload_ranges(config.clone(), cas_client.clone(), original_hash, original_size, inputs) + let result = upload_ranges(config.clone(), cas_client.clone(), original_hash, original_size, inputs, None) .await .unwrap(); let result_hash = MerkleHash::from_hex(result.hash()).unwrap(); @@ -2075,7 +2095,7 @@ mod tests { let expected_next = apply_planned_edits(&expected, &edits); let inputs = edits_to_dirty_inputs(&edits); let result = - upload_ranges(config.clone(), cas_client.clone(), original_hash, original_size, inputs) + upload_ranges(config.clone(), cas_client.clone(), original_hash, original_size, inputs, None) .await .unwrap(); let result_hash = MerkleHash::from_hex(result.hash()).unwrap(); diff --git a/xet_pkg/Cargo.toml b/xet_pkg/Cargo.toml index 25b8eee0b..dd443b4f2 100644 --- a/xet_pkg/Cargo.toml +++ b/xet_pkg/Cargo.toml @@ -85,10 +85,15 @@ tokio_with_wasm = { workspace = true, features = ["rt", "sync"] } [dev-dependencies] async-std = { workspace = true } +chrono = { workspace = true } +ctor = { workspace = true } futures = { workspace = true } serial_test = { workspace = true } smol = { workspace = true } tempfile = { workspace = true } +sha2 = { workspace = true } +dirs = { workspace = true } +uuid = { workspace = true, features = ["v4"] } tokio = { workspace = true, features = [ "rt-multi-thread", "rt", diff --git a/xet_pkg/src/xet_session/mod.rs b/xet_pkg/src/xet_session/mod.rs index 1646824ba..e593cefe2 100644 --- a/xet_pkg/src/xet_session/mod.rs +++ b/xet_pkg/src/xet_session/mod.rs @@ -271,6 +271,8 @@ mod errors; mod file_download_group; #[cfg(not(target_family = "wasm"))] mod file_download_handle; +mod range_upload_commit; +mod range_upload_edit; mod session; mod task_runtime; #[cfg(test)] @@ -287,6 +289,8 @@ pub use file_download_group::{XetDownloadGroupReport, XetFileDownloadGroup, XetF #[cfg(not(target_family = "wasm"))] pub use file_download_handle::{XetDownloadReport, XetFileDownload}; pub use http::{HeaderMap, HeaderValue, header}; +pub use range_upload_commit::{XetRangeUploadCommit, XetRangeUploadCommitBuilder, XetRangeUploadReport}; +pub use range_upload_edit::XetRangeUploadEdit; pub use session::{XetSession, XetSessionBuilder}; pub use task_runtime::XetTaskState; pub use upload_commit::{XetCommitReport, XetFileMetadata, XetUploadCommit, XetUploadCommitBuilder}; diff --git a/xet_pkg/src/xet_session/range_upload_commit.rs b/xet_pkg/src/xet_session/range_upload_commit.rs new file mode 100644 index 000000000..32d8f4f67 --- /dev/null +++ b/xet_pkg/src/xet_session/range_upload_commit.rs @@ -0,0 +1,723 @@ +//! XetRangeUploadCommit — group of edits to an existing file (dirty upload). +//! +//! This is the "dirty upload" layer: instead of uploading an entire file, you +//! specify which byte ranges have changed (dirty ranges) and provide the new +//! data. The untouched regions are pulled from the original file in CAS. +//! +//! ```text +//! with session.new_range_upload(original_hash="abc...", original_size=10000) as commit: +//! commit.edit(1000, 2000).write(b"new data") +//! commit.append(500) +//! report = commit.commit() +//! ``` + +use std::ops::Range; +use std::sync::{Arc, Mutex}; + +use xet_core_structures::merklehash::MerkleHash; +use xet_data::processing::configurations::TranslatorConfig; +use xet_data::processing::{DirtyInput, FileUploadSession, XetFileInfo, create_remote_client}; +use xet_data::progress_tracking::{GroupProgressReport, ItemProgressReport}; +use xet_runtime::utils::UniqueId; + +use super::auth_group_builder::{AuthGroupBuilder, AuthOptions}; +use super::common::create_translator_config; +use super::range_upload_edit::XetRangeUploadEdit; +use super::session::XetSession; +use super::task_runtime::{TaskRuntime, XetTaskState}; +use crate::error::XetError; + +// ── Report ─────────────────────────────────────────────────────────────────── + +/// Report returned by [`XetRangeUploadCommit::commit`]. +#[derive(Clone, Debug)] +#[cfg_attr(feature = "python", pyo3::pyclass(get_all, from_py_object))] +pub struct XetRangeUploadReport { + /// Xet file information for the composed file: hash, size, and optional SHA-256. + pub file_info: XetFileInfo, +} + +// ── Builder ────────────────────────────────────────────────────────────────── + +pub type XetRangeUploadCommitBuilder = AuthGroupBuilder; + +impl AuthGroupBuilder { + /// Create the [`XetRangeUploadCommit`] from an async context. + pub async fn build(self, original_hash: String, original_size: u64) -> Result { + let AuthGroupBuilder { + session, auth_options, .. + } = self; + let parent_runtime = session.inner.task_runtime.clone(); + let child_parent = parent_runtime.clone(); + let commit = parent_runtime + .bridge_async("new_range_upload", async move { + let commit_runtime = child_parent.child()?; + XetRangeUploadCommit::new(session, commit_runtime, auth_options, original_hash, original_size).await + }) + .await?; + Ok(commit) + } + + /// Create the [`XetRangeUploadCommit`] from a sync context. + /// + /// # Errors + /// + /// Returns [`XetError::WrongRuntimeMode`] if the session wraps an external + /// tokio runtime. + /// + /// # Panics + /// + /// Panics if called from within a tokio async runtime on an Owned-mode session. + #[cfg(not(target_family = "wasm"))] + pub fn build_blocking(self, original_hash: String, original_size: u64) -> Result { + let AuthGroupBuilder { + session, auth_options, .. + } = self; + let parent_runtime = session.inner.task_runtime.clone(); + let child_parent = parent_runtime.clone(); + let commit = parent_runtime.bridge_sync("new_range_upload_blocking", async move { + let commit_runtime = child_parent.child()?; + XetRangeUploadCommit::new(session, commit_runtime, auth_options, original_hash, original_size).await + })?; + Ok(commit) + } +} + +// ── XetRangeUploadCommit (public wrapper) ──────────────────────────────────── + +/// API for editing an existing file by uploading only changed byte ranges. +/// +/// Obtain via [`XetSession::new_range_upload`] — configure auth on the returned +/// [`AuthGroupBuilder`], then call [`build`](AuthGroupBuilder::build) (async) or +/// [`build_blocking`](AuthGroupBuilder::build_blocking) (sync). +/// +/// Queue edits with [`edit`](Self::edit), [`insert`](Self::insert), [`delete`](Self::delete), +/// or [`append`](Self::append), then call +/// [`commit`](Self::commit) (async) or [`commit_blocking`](Self::commit_blocking) (sync). +/// +/// This type is cheaply clonable; all clones share the same underlying state. +/// +/// # Errors +/// +/// Returns [`XetError::UserCancelled`] if the parent session has been aborted. +#[derive(Clone)] +pub struct XetRangeUploadCommit { + pub(super) inner: Arc, + pub(super) task_runtime: Arc, +} + +impl XetRangeUploadCommit { + pub(super) async fn new( + session: XetSession, + task_runtime: Arc, + auth_options: AuthOptions, + original_hash: String, + original_size: u64, + ) -> Result { + // Validate auth by creating the translator config (this resolves the endpoint + // and token early, failing fast if auth is invalid). + let config = Arc::new(create_translator_config(&session, auth_options).await?); + let client = create_remote_client(&config, &session.inner.id.to_string(), false).await?; + + // Create upload session for progress tracking + let upload_session = FileUploadSession::new(Arc::clone(&config)).await?; + + let commit_id = UniqueId::new(); + let inner = Arc::new(XetRangeUploadCommitInner { + commit_id, + config, + client, + original_hash, + original_size, + pending_edits: Mutex::new(Vec::new()), + upload_session: Arc::new(std::sync::Mutex::new(Some(upload_session))), + }); + + Ok(Self { inner, task_runtime }) + } + + /// Unique identifier for this commit. + pub fn id(&self) -> UniqueId { + self.inner.commit_id + } + + /// Status of this commit. + pub fn status(&self) -> Result { + self.task_runtime.status() + } + + /// Start a new edit: replace `original_range` with `new_length` bytes. + /// + /// Returns an [`XetRangeUploadEdit`] handle. Feed data incrementally with + /// [`write`](XetRangeUploadEdit::write), then call + /// [`finish`](XetRangeUploadEdit::finish) **before** calling [`commit`]. + pub fn edit(&self, original_range: Range, new_length: u64) -> XetRangeUploadEdit { + let edit = XetRangeUploadEdit::new(original_range, new_length); + self.inner.pending_edits.lock().unwrap().push(edit.clone()); + edit + } + + /// Convenience: insert `new_length` bytes at position `pos`. + /// + /// Equivalent to `edit(pos..pos, new_length)`. + pub fn insert(&self, pos: u64, new_length: u64) -> XetRangeUploadEdit { + self.edit(pos..pos, new_length) + } + + /// Convenience: delete bytes at `start..end`. + /// + /// Equivalent to `edit(start..end, 0)`. + pub fn delete(&self, start: u64, end: u64) -> XetRangeUploadEdit { + self.edit(start..end, 0) + } + + /// Convenience: append `new_length` bytes at the end of the file. + /// + /// Equivalent to `edit(original_size..original_size, new_length)`. + pub fn append(&self, new_length: u64) -> XetRangeUploadEdit { + let original_size = self.inner.original_size; + self.edit(original_size..original_size, new_length) + } + + /// Wait for all edits to be committed and return the result. + pub async fn commit(&self) -> Result { + let inner = Arc::clone(&self.inner); + self.task_runtime + .bridge_async_finalizing("range_upload_commit", false, async move { inner.commit().await }) + .await + } + + /// Blocking version of [`commit`](Self::commit). + /// + /// # Panics + /// + /// Panics if called from within a tokio async runtime. + #[cfg(not(target_family = "wasm"))] + pub fn commit_blocking(&self) -> Result { + let inner = Arc::clone(&self.inner); + self.task_runtime.bridge_sync_finalizing( + "range_upload_commit_blocking", + false, + async move { inner.commit().await }, + ) + } + + /// Cancel all pending edits. + pub fn abort(&self) -> Result<(), XetError> { + let mut pending = self.inner.pending_edits.lock().unwrap(); + pending.clear(); + self.task_runtime.cancel_subtree()?; + Ok(()) + } + + /// Aggregate progress for this commit. + pub fn progress(&self) -> GroupProgressReport { + self.inner + .upload_session + .lock() + .unwrap() + .as_ref() + .map(|s| s.report()) + .unwrap_or_default() + } + + /// Get item reports from the upload session. + pub fn item_reports_from_upload_session(&self) -> std::collections::HashMap { + self.inner + .upload_session + .lock() + .unwrap() + .as_ref() + .map(|s| s.item_reports()) + .unwrap_or_default() + } +} + +// ── XetRangeUploadCommitInner ─────────────────────────────────────────────── + +pub(crate) struct XetRangeUploadCommitInner { + commit_id: UniqueId, + /// Translator config with endpoint, auth, etc. (wrapped in Arc for sharing). + config: Arc, + /// CAS client for fetching original file segments. + client: Arc, + original_hash: String, + original_size: u64, + /// Pending edit handles that will be consumed by commit. + pending_edits: Mutex>, + /// Upload session for progress tracking (created lazily on first edit). + upload_session: Arc>>>, +} + +impl XetRangeUploadCommitInner { + /// Finalise all pending edits and execute the range upload. + async fn commit(self: &Arc) -> Result { + // Finalise each edit and collect DirtyInputs. All edits use **original-file** + // coordinates and must be non-overlapping. The caller is responsible for + // merging any overlapping operations before calling commit(). + let mut dirty_inputs: Vec = { + let mut pending = self.pending_edits.lock().unwrap(); + let mut inputs = Vec::new(); + for edit in pending.drain(..) { + let edit_arc = Arc::new(edit); + let dirty = edit_arc + .finish() + .map_err(|_| XetError::other("edit was already finished before commit"))?; + inputs.push(dirty); + } + // Sort by original_range so the non-overlapping check and merge logic below + // work correctly regardless of edit insertion order. + inputs.sort_by_key(|d| d.original_range.start); + inputs + }; + + tracing::debug!("Committing range upload with {} edits", dirty_inputs.len()); + + // Convert the original hash string to MerkleHash. + let original_hash = MerkleHash::from_hex(&self.original_hash) + .map_err(|e| XetError::other(format!("invalid original_hash: {e}")))?; + + // Run upload_ranges with the config and client. + let upload_session_for_ranges = { + let guard = self.upload_session.lock().unwrap(); + guard.as_ref().cloned() + }; + + let file_info = xet_data::processing::upload_ranges( + Arc::clone(&self.config), + Arc::clone(&self.client), + original_hash, + self.original_size, + dirty_inputs, + upload_session_for_ranges, + ) + .await?; + + Ok(XetRangeUploadReport { file_info }) + } +} + +// ── Tests ──────────────────────────────────────────────────────────────────── + +#[cfg(test)] +mod tests { + use std::sync::OnceLock; + + use http::HeaderMap; + use tempfile::tempdir; + + use super::*; + use crate::xet_session::Sha256Policy; + use crate::xet_session::session::XetSessionBuilder; + + /// Computes the test directory once (date-uuid) and reuses it for all uploads. + static TEST_DIR: OnceLock = OnceLock::new(); + + fn get_test_dir() -> &'static str { + TEST_DIR + .get_or_init(|| { + let date = chrono::Utc::now().format("%Y-%m-%d"); + let uuid = uuid::Uuid::new_v4(); + format!("{date}-{uuid}") + }) + .as_str() + } + + /// Cleanup: remove the test directory and all its contents when the process exits. + #[ctor::ctor(unsafe)] + fn cleanup_test_dir() { + if let Some(dir) = TEST_DIR.get() { + let _ = std::fs::remove_dir_all(dir); + } + } + + /// Helper: read the HF Hub token, preferring the HF_TOKEN env var and falling back + /// to the default cache path. + fn read_hf_token() -> String { + if let Ok(token) = std::env::var("HF_TOKEN") { + let token = token.trim().to_string(); + if !token.is_empty() { + return token; + } + } + let token_path = dirs::home_dir() + .map(|d| d.join(".cache/huggingface/token")) + .expect("could not resolve home dir"); + std::fs::read_to_string(token_path) + .expect("failed to read HF token") + .trim() + .to_string() + } + + fn upload_file(session: &XetSession, endpoint: &str, data: &[u8], name: &str) -> XetFileInfo { + let dir = get_test_dir(); + let full_name = format!("{}/{}", dir, name); + let commit = session + .new_upload_commit() + .unwrap() + .with_endpoint(endpoint) + .build_blocking() + .unwrap(); + let _handle = commit + .upload_bytes_blocking(data.to_vec(), Sha256Policy::Compute, Some(full_name)) + .unwrap(); + let results = commit.commit_blocking().unwrap(); + let meta = results.uploads.into_values().next().expect("one uploaded file"); + meta.xet_info.clone() + } + + #[test] + fn test_range_upload_edit_basic() { + let temp = tempdir().unwrap(); + let endpoint = format!("local://{}", temp.path().join("cas").display()); + let session = XetSessionBuilder::new().build().unwrap(); + + // Upload an original file (13 bytes: "Hello, World!") + let original_data = b"Hello, World!"; + let original_info = upload_file(&session, &endpoint, original_data, "original.bin"); + + // Create a range upload commit + let commit = session + .new_range_upload() + .unwrap() + .with_endpoint(&endpoint) + .build_blocking(original_info.hash, original_info.file_size.unwrap()) + .unwrap(); + + // Edit: replace bytes 7..12 (5 bytes: ", Wor") with new data (8 bytes: "Universe") + // Expected file size: 13 - 5 + 8 = 16 + let edit = commit.edit(7..12, 8); + edit.write(b"Universe"); + + // Commit + let report = commit.commit_blocking().unwrap(); + assert_eq!(report.file_info.file_size, Some(16)); + } + + #[test] + fn test_range_upload_insert() { + let temp = tempdir().unwrap(); + let endpoint = format!("local://{}", temp.path().join("cas").display()); + let session = XetSessionBuilder::new().build().unwrap(); + + let original_data = b"Hello World!"; + let original_info = upload_file(&session, &endpoint, original_data, "original.bin"); + + let commit = session + .new_range_upload() + .unwrap() + .with_endpoint(&endpoint) + .build_blocking(original_info.hash, original_info.file_size.unwrap()) + .unwrap(); + + // Insert 7 bytes at position 5 (empty original range, so new_length = 7) + let edit = commit.insert(5, 7); + edit.write(b" Beautiful"); + + let report = commit.commit_blocking().unwrap(); + // Original 12 bytes + 7 inserted = 19 bytes + assert_eq!(report.file_info.file_size, Some(19)); + } + + #[test] + fn test_range_upload_delete() { + let temp = tempdir().unwrap(); + let endpoint = format!("local://{}", temp.path().join("cas").display()); + let session = XetSessionBuilder::new().build().unwrap(); + + let original_data = b"Hello, World!"; + let original_info = upload_file(&session, &endpoint, original_data, "original.bin"); + + let commit = session + .new_range_upload() + .unwrap() + .with_endpoint(&endpoint) + .build_blocking(original_info.hash, original_info.file_size.unwrap()) + .unwrap(); + + // Delete bytes 7..12 (5 bytes: ", Wor") + // 13 - 5 = 8 bytes + let _edit = commit.delete(7, 12); + + let report = commit.commit_blocking().unwrap(); + assert_eq!(report.file_info.file_size, Some(8)); + } + + #[test] + fn test_range_upload_append() { + let temp = tempdir().unwrap(); + let endpoint = format!("local://{}", temp.path().join("cas").display()); + let session = XetSessionBuilder::new().build().unwrap(); + + let original_data = b"Hello"; + let original_info = upload_file(&session, &endpoint, original_data, "original.bin"); + + let commit = session + .new_range_upload() + .unwrap() + .with_endpoint(&endpoint) + .build_blocking(original_info.hash, original_info.file_size.unwrap()) + .unwrap(); + + // Append 6 bytes + let edit = commit.append(6); + edit.write(b" World"); + + let report = commit.commit_blocking().unwrap(); + assert_eq!(report.file_info.file_size, Some(11)); + } + + #[test] + fn test_range_upload_multiple_edits() { + let temp = tempdir().unwrap(); + let endpoint = format!("local://{}", temp.path().join("cas").display()); + let session = XetSessionBuilder::new().build().unwrap(); + + let original_data = b"Hello, World!"; + let original_info = upload_file(&session, &endpoint, original_data, "original.bin"); + + let commit = session + .new_range_upload() + .unwrap() + .with_endpoint(&endpoint) + .build_blocking(original_info.hash, original_info.file_size.unwrap()) + .unwrap(); + + // Multiple edits + // 7..12 (5 bytes) -> 8 bytes (Universe) => +3 + // 12..12 (0 bytes) -> 1 byte (!) => +1 + // Total: 13 + 3 + 1 = 17 + commit.edit(7..12, 8).write(b"Universe"); + commit.edit(12..12, 1).write(b"!"); + + let report = commit.commit_blocking().unwrap(); + assert_eq!(report.file_info.file_size, Some(17)); + } + + // ── E2E tests against HuggingFace Hub ───────────────────────────────────── + + /// Helper: upload a single file to the HF Hub repo. + fn upload_to_hub(session: &XetSession, data: &[u8], name: &str) -> XetFileInfo { + let dir = get_test_dir(); + let full_name = format!("{}/{}", dir, name); + let token = read_hf_token(); + let mut headers = HeaderMap::new(); + headers.insert( + http::header::AUTHORIZATION, + http::header::HeaderValue::from_str(&format!("Bearer {}", token)).unwrap(), + ); + let refresh_url = + "https://huggingface.co/api/buckets/hf-internal-testing/test-xet-core/xet-write-token".to_string(); + + let commit = session + .new_upload_commit() + .unwrap() + .with_token_refresh_url(refresh_url, headers) + .build_blocking() + .unwrap(); + + let _handle = commit + .upload_bytes_blocking(data.to_vec(), Sha256Policy::Compute, Some(full_name)) + .unwrap(); + let results = commit.commit_blocking().unwrap(); + let meta = results.uploads.into_values().next().expect("one uploaded file"); + meta.xet_info.clone() + } + + fn make_auth_headers() -> HeaderMap { + let token = read_hf_token(); + let mut headers = HeaderMap::new(); + headers.insert( + http::header::AUTHORIZATION, + http::header::HeaderValue::from_str(&format!("Bearer {}", token)).unwrap(), + ); + headers + } + + fn make_write_refresh_url() -> String { + "https://huggingface.co/api/buckets/hf-internal-testing/test-xet-core/xet-write-token".to_string() + } + + fn make_read_refresh_url() -> String { + "https://huggingface.co/api/buckets/hf-internal-testing/test-xet-core/xet-read-token".to_string() + } + + #[test] + fn test_e2e_range_upload_hub() { + futures::executor::block_on(async { + let session = XetSessionBuilder::new().build().unwrap(); + + // ── Step 1: Create and upload an original file to HF Hub ────────────── + let original_data = b"Hello, World! This is a test file for range upload."; + let original_info = upload_to_hub(&session, original_data, "original.txt"); + println!("Original: hash={}, size={}", original_info.hash, original_info.file_size.unwrap()); + + // Verify the hash matches + use sha2::{Digest, Sha256}; + let hash_bytes: Vec = Sha256::digest(original_data).to_vec(); + let expected_sha: String = hash_bytes.iter().map(|b| format!("{:02x}", b)).collect(); + assert_eq!(original_info.sha256.as_deref(), Some(expected_sha.as_str())); + + // ── Step 2: Download the original file to verify contents ────────────── + let dl_session = XetSessionBuilder::new().build().unwrap(); + + let group = dl_session + .new_file_download_group() + .unwrap() + .with_token_refresh_url(make_read_refresh_url(), make_auth_headers()) + .build_blocking() + .unwrap(); + + let dest_path = tempfile::tempdir().unwrap().path().join("downloaded.txt"); + let _handle = group + .download_file_to_path_blocking(original_info.clone(), dest_path.clone()) + .unwrap(); + + let report = group.finish_blocking().unwrap(); + assert_eq!(report.downloads.len(), 1); + + let downloaded_data = std::fs::read(&dest_path).unwrap(); + assert_eq!(downloaded_data, original_data); + + // ── Step 3: Perform a range upload (edit) ───────────────────────────── + let edit_data = b"Universe! "; + let write_headers = make_auth_headers(); + + let commit = session + .new_range_upload() + .unwrap() + .with_token_refresh_url(make_write_refresh_url(), write_headers) + .build_blocking(original_info.hash, original_info.file_size.unwrap()) + .unwrap(); + + // Edit: replace bytes 0..13 ("Hello, World!") with "Universe! " (10 bytes) + // Original: 51 bytes ("Hello, World! This is a test file for range upload.") + // After: 51 - 13 + 10 = 48 bytes + + let edit = commit.edit(0..13, 10); + edit.write(edit_data); + + let report = commit.commit_blocking().unwrap(); + assert_eq!(report.file_info.file_size, Some(48)); + + // ── Step 4: Download the modified file and verify ────────────────────── + let dl_session2 = XetSessionBuilder::new().build().unwrap(); + let dl_headers2 = make_auth_headers(); + + let group2 = dl_session2 + .new_file_download_group() + .unwrap() + .with_token_refresh_url(make_read_refresh_url(), dl_headers2) + .build_blocking() + .unwrap(); + + let dest_path2 = tempfile::tempdir().unwrap().path().join("downloaded2.txt"); + let _handle2 = group2 + .download_file_to_path_blocking(report.file_info.clone(), dest_path2.clone()) + .unwrap(); + + let report2 = group2.finish_blocking().unwrap(); + assert_eq!(report2.downloads.len(), 1); + + let modified_data = std::fs::read(&dest_path2).unwrap(); + let expected_modified = b"Universe! This is a test file for range upload."; + assert_eq!(expected_modified.len(), 48); + assert_eq!(modified_data, expected_modified); + }); + } + + #[test] + fn test_e2e_range_upload_insert_hub() { + futures::executor::block_on(async { + let session = XetSessionBuilder::new().build().unwrap(); + + // Upload original + let original_data = b"ABCDEF"; + let original_info = upload_to_hub(&session, original_data, "insert_test.txt"); + + // Insert 3 bytes at position 2: "XYZ" + // Result: "ABXYZCDEF" (9 bytes) + let write_headers = make_auth_headers(); + + let commit = session + .new_range_upload() + .unwrap() + .with_token_refresh_url(make_write_refresh_url(), write_headers) + .build_blocking(original_info.hash, original_info.file_size.unwrap()) + .unwrap(); + + commit.insert(2, 3).write(b"XYZ"); + + let report = commit.commit_blocking().unwrap(); + assert_eq!(report.file_info.file_size, Some(9)); + + // Verify by downloading + let dl_session = XetSessionBuilder::new().build().unwrap(); + let dl_headers = make_auth_headers(); + + let group = dl_session + .new_file_download_group() + .unwrap() + .with_token_refresh_url(make_read_refresh_url(), dl_headers) + .build_blocking() + .unwrap(); + + let dest = tempfile::tempdir().unwrap().path().join("inserted.txt"); + group + .download_file_to_path_blocking(report.file_info.clone(), dest.clone()) + .unwrap(); + + group.finish_blocking().unwrap(); + + let data = std::fs::read(&dest).unwrap(); + assert_eq!(data, b"ABXYZCDEF"); + }); + } + + #[test] + fn test_e2e_range_upload_delete_hub() { + futures::executor::block_on(async { + let session = XetSessionBuilder::new().build().unwrap(); + + let original_data = b"Hello, World!"; + let original_info = upload_to_hub(&session, original_data, "delete_test.txt"); + + // Delete bytes 5..12 (", World") => 7 bytes removed + // Result: "Hello!" (6 bytes) + let write_headers = make_auth_headers(); + + let commit = session + .new_range_upload() + .unwrap() + .with_token_refresh_url(make_write_refresh_url(), write_headers) + .build_blocking(original_info.hash, original_info.file_size.unwrap()) + .unwrap(); + + commit.delete(5, 12); + + let report = commit.commit_blocking().unwrap(); + assert_eq!(report.file_info.file_size, Some(6)); + + // Verify + let dl_session = XetSessionBuilder::new().build().unwrap(); + let dl_headers = make_auth_headers(); + + let group = dl_session + .new_file_download_group() + .unwrap() + .with_token_refresh_url(make_read_refresh_url(), dl_headers) + .build_blocking() + .unwrap(); + + let dest = tempfile::tempdir().unwrap().path().join("deleted.txt"); + group + .download_file_to_path_blocking(report.file_info.clone(), dest.clone()) + .unwrap(); + + group.finish_blocking().unwrap(); + + let data = std::fs::read(&dest).unwrap(); + assert_eq!(data, b"Hello!"); + }); + } +} diff --git a/xet_pkg/src/xet_session/range_upload_commit_tests.rs b/xet_pkg/src/xet_session/range_upload_commit_tests.rs new file mode 100644 index 000000000..31ea282f0 --- /dev/null +++ b/xet_pkg/src/xet_session/range_upload_commit_tests.rs @@ -0,0 +1,131 @@ +//! Tests for XetRangeUploadCommit. + +use std::ops::Range; +use std::sync::Arc; + +use tempfile::tempdir; +use xet_data::processing::{Sha256Policy, XetFileInfo, FileUploadSession}; +use xet_runtime::core::XetContext; + +use super::super::session::XetSessionBuilder; +use super::*; + +async fn upload_file(session: &XetSession, endpoint: &str, data: &[u8], name: &str) -> XetFileInfo { + let commit = session + .new_upload_commit() + .unwrap() + .with_endpoint(endpoint) + .build_blocking() + .unwrap(); + let _handle = commit + .upload_bytes(data.to_vec(), Sha256Policy::Compute, Some(name.into())) + .unwrap(); + let results = commit.commit_blocking().unwrap(); + let meta = results.uploads.into_values().next().expect("one uploaded file"); + meta.xet_info.clone() +} + +#[test] +fn test_range_upload_commit_basic() -> anyhow::Result<()> { + let temp = tempdir()?; + let cas_path = temp.path().join("cas"); + let endpoint = format!("local://{}", cas_path.display()); + let session = XetSessionBuilder::new().build()?; + + // Upload an original file + let original_data = b"Hello, World! This is the original content."; + let original_info = upload_file(&session, &endpoint, original_data, "original.bin")?; + + // Create a range upload commit + let commit = session + .new_range_upload()? + .with_endpoint(&endpoint) + .build_blocking(original_info.hash, original_info.file_size.unwrap())?; + + // Edit: replace bytes 7..12 with new data + let edit = commit.edit(7..12, 10); + edit.write(b"Universe"); + commit.commit_blocking()?; + + // Verify: the composed file should have "Hello, Universe! This is the original content." + // But since we don't have SHA-256, we just verify the hash changed + let report = commit.commit_blocking()?; + + assert_eq!( + report.file_info.file_size, + Some(original_info.file_size.unwrap()) + ); + + Ok(()) +} + +#[test] +fn test_range_upload_commit_insert() -> anyhow::Result<()> { + let temp = tempdir()?; + let endpoint = format!("local://{}", temp.path().join("cas").display()); + let session = XetSessionBuilder::new().build()?; + + let original_data = b"Hello World!"; + let original_info = upload_file(&session, &endpoint, original_data, "original.bin")?; + + let commit = session + .new_range_upload()? + .with_endpoint(&endpoint) + .build_blocking(original_info.hash, original_info.file_size.unwrap())?; + + // Insert 7 bytes at position 5 + let edit = commit.insert(5, 7); + edit.write(b" Beautiful"); + + let report = commit.commit_blocking()?; + assert_eq!(report.file_info.file_size, Some(19)); + + Ok(()) +} + +#[test] +fn test_range_upload_commit_delete() -> anyhow::Result<()> { + let temp = tempdir()?; + let endpoint = format!("local://{}", temp.path().join("cas").display()); + let session = XetSessionBuilder::new().build()?; + + let original_data = b"Hello, World!"; + let original_info = upload_file(&session, &endpoint, original_data, "original.bin")?; + + let commit = session + .new_range_upload()? + .with_endpoint(&endpoint) + .build_blocking(original_info.hash, original_info.file_size.unwrap())?; + + // Delete bytes 7..12 (", Wor") + let edit = commit.delete(7, 12); + + let report = commit.commit_blocking()?; + assert_eq!(report.file_info.file_size, Some(6)); + + Ok(()) +} + +#[test] +fn test_range_upload_commit_append() -> anyhow::Result<()> { + let temp = tempdir()?; + let endpoint = format!("local://{}", temp.path().join("cas").display()); + let session = XetSessionBuilder::new().build()?; + + let original_data = b"Hello"; + let original_info = upload_file(&session, &endpoint, original_data, "original.bin")?; + + let commit = session + .new_range_upload()? + .with_endpoint(&endpoint) + .build_blocking(original_info.hash, original_info.file_size.unwrap())?; + + // Append 6 bytes + let edit = commit.append(6); + edit.write(b" World"); + + let report = commit.commit_blocking()?; + assert_eq!(report.file_info.file_size, Some(11)); + + Ok(()) +} \ No newline at end of file diff --git a/xet_pkg/src/xet_session/range_upload_edit.rs b/xet_pkg/src/xet_session/range_upload_edit.rs new file mode 100644 index 000000000..6771f81de --- /dev/null +++ b/xet_pkg/src/xet_session/range_upload_edit.rs @@ -0,0 +1,126 @@ +//! XetRangeUploadEdit — pending data for a single edit within a range upload. + +use std::ops::Range; +use std::sync::{Arc, Mutex}; + +use xet_data::processing::DirtyInput; + +#[derive(Debug)] +pub enum RangeUploadEditError { + AlreadyFinished, +} + +impl std::fmt::Display for RangeUploadEditError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + RangeUploadEditError::AlreadyFinished => write!(f, "edit already finished"), + } + } +} + +impl std::error::Error for RangeUploadEditError {} + +// ── XetRangeUploadEditInner ───────────────────────────────────────────────── + +pub(super) struct XetRangeUploadEditInner { + pub(super) original_range: Range, + pub(super) new_length: u64, + /// The accumulated data for this edit. None means it has been consumed. + pub(super) data: Mutex>>, +} + +impl XetRangeUploadEditInner { + fn write(&self, data: &[u8]) { + let mut guard = self.data.lock().unwrap(); + if let Some(buf) = guard.as_mut() { + buf.extend_from_slice(data); + } + } + + /// Finalise the edit, returning the pending [`DirtyInput`] and clearing the buffer. + fn finish(self: &Arc) -> Result { + let mut guard = self.data.lock().unwrap(); + let data = guard.take().ok_or(RangeUploadEditError::AlreadyFinished)?; + Ok(DirtyInput { + original_range: self.original_range.clone(), + new_length: self.new_length, + reader: Box::pin(std::io::Cursor::new(data.to_vec())), + }) + } + + /// Returns the pending data without blocking, or `None` if already finished. + fn try_finish(self: &Arc) -> Option { + let mut guard = self.data.lock().ok()?; + let data = guard.take()?; + Some(DirtyInput { + original_range: self.original_range.clone(), + new_length: self.new_length, + reader: Box::pin(std::io::Cursor::new(data.to_vec())), + }) + } +} + +// ── XetRangeUploadEdit (public wrapper) ────────────────────────────────────── + +/// Handle for a single edit within a [`XetRangeUploadCommit`]. +/// +/// Returned by [`XetRangeUploadCommit::edit`], [`insert`], and [`delete`]. +/// Feed data incrementally with [`write`], then call [`finish`] to obtain the +/// pending [`DirtyInput`]. +/// +/// **`finish` must be called before [`XetRangeUploadCommit::commit`]**. +/// +/// [`write`]: Self::write +/// [`finish`]: Self::finish +/// [`insert`]: XetRangeUploadCommit::insert +/// [`delete`]: XetRangeUploadCommit::delete +#[derive(Clone)] +pub struct XetRangeUploadEdit { + pub(super) inner: Arc, +} + +impl std::fmt::Debug for XetRangeUploadEdit { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("XetRangeUploadEdit") + .field("original_range", &self.inner.original_range) + .field("new_length", &self.inner.new_length) + .finish_non_exhaustive() + } +} + +impl XetRangeUploadEdit { + pub(super) fn new(original_range: Range, new_length: u64) -> Self { + Self { + inner: Arc::new(XetRangeUploadEditInner { + original_range, + new_length, + data: Mutex::new(Some(Vec::new())), + }), + } + } + + /// Feed data into this edit. + /// + /// May be called any number of times before [`finish`]. + /// + /// [`finish`]: Self::finish + pub fn write(&self, data: &[u8]) { + self.inner.write(data); + } + + /// Finalise the edit, returning the pending [`DirtyInput`]. + /// + /// Must be called before [`XetRangeUploadCommit::commit`]. A second call returns `Err` + /// after a successful finish; use [`try_finish`] to read cached data without + /// finalising again. + /// + /// [`try_finish`]: Self::try_finish + pub fn finish(self: &Arc) -> Result { + self.inner.finish() + } + + /// Returns the pending data without blocking, or `None` if already finished. + pub fn try_finish(self: &Arc) -> Option { + self.inner.try_finish() + } +} diff --git a/xet_pkg/src/xet_session/session.rs b/xet_pkg/src/xet_session/session.rs index 95d661cfa..a623f0648 100644 --- a/xet_pkg/src/xet_session/session.rs +++ b/xet_pkg/src/xet_session/session.rs @@ -17,6 +17,7 @@ use super::download_stream_group::{ use super::errors::SessionError; #[cfg(not(target_family = "wasm"))] use super::file_download_group::XetFileDownloadGroupBuilder; +use super::range_upload_commit::XetRangeUploadCommitBuilder; use super::task_runtime::{TaskRuntime, XetTaskState}; use super::upload_commit::XetUploadCommitBuilder; @@ -313,6 +314,35 @@ impl XetSession { Ok(XetDownloadStreamGroupBuilder::new(self.clone())) } + /// Create a [`XetRangeUploadCommitBuilder`] for configuring and constructing a range upload (dirty upload). + /// + /// A range upload edits an existing file by uploading only the changed byte ranges. + /// The untouched regions are pulled from the original file in CAS. + /// + /// Configure the builder with any combination of: + /// - [`with_endpoint`](crate::xet_session::auth_group_builder::AuthGroupBuilder::with_endpoint) — CAS server URL + /// - [`with_token_info`](crate::xet_session::auth_group_builder::AuthGroupBuilder::with_token_info) — CAS token and + /// expiry + /// - [`with_token_refresh_url`](crate::xet_session::auth_group_builder::AuthGroupBuilder::with_token_refresh_url) — + /// URL to refresh the token + /// + /// Then call [`build`](XetRangeUploadCommitBuilder::build) (async) or + /// [`build_blocking`](XetRangeUploadCommitBuilder::build_blocking) (sync) with the original + /// file's hash and size. Queue edits with [`edit`](XetRangeUploadCommit::edit), + /// [`insert`](XetRangeUploadCommit::insert), [`delete`](XetRangeUploadCommit::delete), or + /// [`append`](XetRangeUploadCommit::append), then call + /// [`commit`](XetRangeUploadCommit::commit) (async) or + /// [`commit_blocking`](XetRangeUploadCommit::commit_blocking) (sync). + /// + /// Returns `Err(SessionError::UserCancelled)` if the session has been aborted. + #[cfg(not(target_family = "wasm"))] + pub fn new_range_upload(&self) -> Result { + self.inner.task_runtime.check_state("new_range_upload")?; + #[cfg(feature = "fd-track")] + report_fd_count("XetSession::new_range_upload"); + Ok(XetRangeUploadCommitBuilder::new(self.clone())) + } + pub fn status(&self) -> Result { self.inner.task_runtime.status() }