update truss transfer (#1383)

* update truss transfer * update readme * add sync all command
basetenlabs · Feb 11, 2025 · 4df752d · 4df752d
1 parent 35e8d3e
commit 4df752d
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 8 deletions.
diff --git a/.github/workflows/truss-transfer-test.yml b/.github/workflows/truss-transfer-test.yml
@@ -3,7 +3,7 @@ name: Test truss-transfer
 on:
   push:
     tags:
-      - '*'
+      - "*"
   pull_request:
   workflow_dispatch:
 
@@ -12,7 +12,7 @@ permissions:
 
 defaults:
   run:
-    working-directory: truss-transfer  # Keeps default for run commands
+    working-directory: truss-transfer # Keeps default for run commands
 
 jobs:
   test-truss-transfer:
@@ -21,15 +21,15 @@ jobs:
       - uses: actions/checkout@v4
       - uses: actions/setup-python@v5
         with:
-          python-version: '3.10'
+          python-version: "3.10"
       - name: Create venv
         run: python3 -m venv .venv
-      - name: Build wheels
+      - name: Build wheels for python
         uses: PyO3/maturin-action@v1
         with:
           working-directory: truss-transfer
           command: develop
-          sccache: 'true'
+          sccache: "true"
       - name: Try import truss-transfer-maturin-publish
         run: |
           source .venv/bin/activate

diff --git a/truss-transfer/README.md b/truss-transfer/README.md
@@ -24,6 +24,10 @@ def lazy_data_loader(download_dir: str):
 
 Compiling the libary as musl-linux target for cross-platform usage.
 ```
+# Add one-time installations
+# apt-get install -y musl-tools libssl-dev libatomic-ops-dev
+# rustup target add x86_64-unknown-linux-musl
+
 # To build with cargo:
 cargo build --release --target x86_64-unknown-linux-musl --features cli --bin truss_transfer_cli
 ```

diff --git a/truss-transfer/src/lib.rs b/truss-transfer/src/lib.rs
@@ -276,9 +276,17 @@ async fn download_file_with_cache(
             download_to_path(client, url, &destination, size).await?;
         } else {
             // success in caching => symlink to final dest
-            println!("[INFO] Download to b10cache successful, creating symlink to final destination.");
+            println!(
+                "[INFO] Download to b10cache successful, creating symlink to final destination."
+            );
             if let Err(e) = create_symlink_or_skip(&cache_path, &destination) {
-                println!("[DEBUG] Symlink to data dir failed: {e}");
+                println!("[WARN] Symlink failed: {e}. Falling back to direct download.");
+                if let Err(download_err) = download_to_path(client, url, &destination, size).await {
+                    println!("[ERROR] Direct download failed: {download_err}");
+                    return Err(anyhow!(
+                        "Failed to create symlink and direct download also failed"
+                    ));
+                }
             }
         }
     } else {
@@ -316,6 +324,9 @@ async fn download_to_path(client: &Client, url: &str, path: &Path, size: i64) ->
         file.write_all(&chunk).await?;
     }
 
+    // Ensure data is flushed to disk.
+    file.sync_all().await?;
+
     // Optional size check
     if size > 0 {
         let written = file.metadata().await?.len();
@@ -324,8 +335,12 @@ async fn download_to_path(client: &Client, url: &str, path: &Path, size: i64) ->
                 "Warning: downloaded file size mismatch (expected {}, got {}) at {:?}",
                 size, written, path
             );
+            // TODO: fail if size has large discrepancy, e.g. > 10%
         } else {
-            println!("[INFO] Download size matches expected size: {size} bytes.");
+            println!(
+                "[INFO] Download size of {:?} matches expected size of {}",
+                path, size
+            );
         }
     }