Merge pull request #51 from imbue-ai/jacob/reliability-fix

Fix race condition in parallel sandbox downloads

Author: jsk11235

scripts/modal_sandbox.py

@@ -77,59 +77,22 @@ def copy_dir_to_sandbox(sandbox, local_dir: str, remote_dir: str) -> None:
 
 
 def copy_from_sandbox(sandbox, remote_path: str, local_path: str) -> None:
-    """Copy a file or directory from the sandbox to local filesystem using tar."""
+    """Copy a file from the sandbox to local filesystem."""
     logger.info("Downloading %s to %s...", remote_path, local_path)
 
-    # Use tar on the sandbox to create an archive of the remote path
-    # This handles both files and directories uniformly
-    tar_remote_path = "/tmp/.download_transfer.tar"
+    # Read file content directly from sandbox
+    with sandbox.open(remote_path, "rb") as f:
+        data = f.read()
 
-    # Create tar archive on sandbox
-    # Use -C to change to parent directory and archive just the basename
-    # This preserves the directory structure correctly
-    remote_parent = os.path.dirname(remote_path.rstrip("/")) or "/"
-    remote_basename = os.path.basename(remote_path.rstrip("/"))
-
-    logger.info("Creating tar archive on sandbox...")
-    process = sandbox.exec(
-        "tar", "-cf", tar_remote_path, "-C", remote_parent, remote_basename
-    )
-    process.wait()
-    if process.returncode != 0:
-        stderr = process.stderr.read()
-        raise click.ClickException(f"Failed to create tar archive on sandbox: {stderr}")
-
-    # Read the tar archive from sandbox
-    logger.info("Transferring tar archive from sandbox...")
-    with sandbox.open(tar_remote_path, "rb") as f:
-        tar_data = f.read()
-
-    logger.info("Received tar archive (%d bytes)", len(tar_data))
-
-    # Clean up tar file on sandbox
-    sandbox.exec("rm", "-f", tar_remote_path).wait()
-
-    # Extract tar archive locally
-    tar_buffer = io.BytesIO(tar_data)
+    logger.info("Received %d bytes", len(data))
 
     # Create parent directory if needed
     local_parent = os.path.dirname(local_path.rstrip("/")) or "."
     os.makedirs(local_parent, exist_ok=True)
 
-    logger.info("Extracting tar archive to %s...", local_parent)
-    with tarfile.open(fileobj=tar_buffer, mode="r") as tar:
-        tar.extractall(path=local_parent)
-
-    # If local_path differs from the extracted name, rename
-    extracted_path = os.path.join(local_parent, remote_basename)
-    if extracted_path != local_path and os.path.exists(extracted_path):
-        if os.path.exists(local_path):
-            # Remove existing target to allow rename
-            if os.path.isdir(local_path):
-                shutil.rmtree(local_path)
-            else:
-                os.remove(local_path)
-        os.rename(extracted_path, local_path)
+    # Write to local file
+    with open(local_path, "wb") as f:
+        f.write(data)
 
     logger.info("Download complete: %s -> %s", remote_path, local_path)
 
@@ -353,12 +316,11 @@ def download(sandbox_id: str, paths: tuple[str, ...]):
     SANDBOX_ID is the Modal sandbox ID to download from.
 
     PATHS are one or more path specifications in the format "remote_path:local_path".
-    Each specification downloads the remote path to the local path.
-    Both files and directories are supported.
+    Each specification downloads the remote file to the local path.
 
     Examples:
 
-        modal_sandbox.py download sb-abc123 "/app/results:/tmp/results"
+        modal_sandbox.py download sb-abc123 "/tmp/junit.xml:./results/junit.xml"
 
         modal_sandbox.py download sb-abc123 "/app/out:./out" "/app/logs:./logs"
     """

src/provider/default.rs

@@ -621,4 +621,97 @@ mod tests {
             "arg with space should be quoted: {result}"
         );
     }
+
+    /// Integration test for Modal sandbox download functionality via DefaultProvider.
+    ///
+    /// This test requires Modal credentials (MODAL_TOKEN_ID and MODAL_TOKEN_SECRET).
+    /// Skips automatically if credentials are not present.
+    #[tokio::test]
+    async fn modal_download_junit_xml() -> Result<(), Box<dyn std::error::Error>> {
+        use crate::config::{DefaultProviderConfig, SandboxConfig};
+        use crate::provider::SandboxProvider;
+        use futures::StreamExt;
+
+        // Skip if Modal credentials are not available
+        if std::env::var("MODAL_TOKEN_ID").is_err() || std::env::var("MODAL_TOKEN_SECRET").is_err()
+        {
+            eprintln!(
+                "Skipping modal_download_junit_xml: MODAL_TOKEN_ID or MODAL_TOKEN_SECRET not set"
+            );
+            return Ok(());
+        }
+
+        // Create a temp dir with a minimal Dockerfile
+        let temp_dir = tempfile::tempdir()?;
+        let dockerfile_path = temp_dir.path().join("Dockerfile.test");
+        std::fs::write(&dockerfile_path, "FROM python:3.11-slim\n")?;
+
+        // Configure DefaultProvider to use modal_sandbox.py
+        // Use @modal_sandbox.py notation which resolves to the bundled script
+        let config = DefaultProviderConfig {
+            prepare_command: Some("uv run @modal_sandbox.py prepare Dockerfile.test".to_string()),
+            create_command: "uv run @modal_sandbox.py create {image_id}".to_string(),
+            exec_command: "uv run @modal_sandbox.py exec {sandbox_id} {command}".to_string(),
+            destroy_command: "uv run @modal_sandbox.py destroy {sandbox_id}".to_string(),
+            download_command: Some(
+                "uv run @modal_sandbox.py download {sandbox_id} {paths}".to_string(),
+            ),
+            working_dir: Some(temp_dir.path().to_path_buf()),
+            timeout_secs: 300,
+            env: Default::default(),
+            copy_dirs: vec![],
+        };
+
+        let provider = DefaultProvider::from_config(config, &[]).await?;
+
+        // Create sandbox
+        let sandbox_config = SandboxConfig {
+            id: "test-download".to_string(),
+            working_dir: None,
+            env: vec![],
+            copy_dirs: vec![],
+        };
+
+        let sandbox = provider.create_sandbox(&sandbox_config).await?;
+
+        // Write test junit.xml content
+        let test_content = r#"<?xml version="1.0" encoding="UTF-8"?>
+<testsuites>
+  <testsuite name="test_suite" tests="2" failures="0">
+    <testcase name="test_one" classname="tests.test_example" time="0.001"/>
+    <testcase name="test_two" classname="tests.test_example" time="0.002"/>
+  </testsuite>
+</testsuites>"#;
+
+        // Write the file to sandbox using exec
+        let write_cmd = Command::new("sh").arg("-c").arg(format!(
+            "cat > /tmp/junit.xml << 'EOF'\n{}\nEOF",
+            test_content
+        ));
+
+        let mut stream = sandbox.exec_stream(&write_cmd).await?;
+        while stream.next().await.is_some() {}
+
+        // Download the file
+        let download_dir = tempfile::tempdir()?;
+        let local_file = download_dir.path().join("downloaded.xml");
+        let remote_path = std::path::Path::new("/tmp/junit.xml");
+
+        sandbox
+            .download(&[(remote_path, local_file.as_path())])
+            .await?;
+
+        // Verify content
+        let downloaded = std::fs::read_to_string(&local_file)?;
+        assert_eq!(
+            test_content.trim(),
+            downloaded.trim(),
+            "Downloaded content should match"
+        );
+
+        // Cleanup
+        sandbox.terminate().await?;
+
+        Ok(())
+    }
 }