Skip to content

Add sandbox retry logic #345

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
mendral-app wants to merge 1 commit into from
Closed

Add sandbox retry logic #345

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
94 changes: 70 additions & 24 deletions cli/drive.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import (
"net/url"
"os"
"strings"
"time"

blaxel "github.com/blaxel-ai/sdk-go"
"github.com/blaxel-ai/toolkit/cli/core"
Expand Down Expand Up @@ -72,6 +73,17 @@ type sandboxAPIError struct {
Error string `json:"error"`
}

// sandboxErrorResponse represents the full structured error response from the sandbox API,
// including retry guidance.
type sandboxErrorResponse struct {
Error struct {
Code string `json:"code"`
Message string `json:"message"`
Retryable bool `json:"retryable"`
Action string `json:"action"`
} `json:"error"`
}

func DriveCmd() *cobra.Command {
cmd := &cobra.Command{
Use: "drive",
Expand Down Expand Up @@ -170,18 +182,11 @@ the blfs filesystem. It can be used as a recovery tool when mounts are lost.`,
core.ExitWithError(err)
}

resp, err := sandboxRequest(ctx, http.MethodPost, sandboxURL, "/drives/mount", token, bytes.NewReader(jsonBody))
resp, respBody, err := sandboxRequestWithRetry(ctx, http.MethodPost, sandboxURL, "/drives/mount", token, jsonBody)
if err != nil {
core.PrintError("Drive mount", err)
core.ExitWithError(err)
}
defer resp.Body.Close()

respBody, err := io.ReadAll(resp.Body)
if err != nil {
core.PrintError("Drive mount", fmt.Errorf("failed to read response: %w", err))
core.ExitWithError(err)
}

if resp.StatusCode != http.StatusOK {
handleSandboxAPIError(respBody, resp.StatusCode, "mount drive")
Expand Down Expand Up @@ -238,18 +243,11 @@ func DriveUnmountCmd() *cobra.Command {
encodedPath := url.PathEscape(strings.TrimPrefix(mountPath, "/"))
apiPath := fmt.Sprintf("/drives/mount/%s", encodedPath)

resp, err := sandboxRequest(ctx, http.MethodDelete, sandboxURL, apiPath, token, nil)
resp, respBody, err := sandboxRequestWithRetry(ctx, http.MethodDelete, sandboxURL, apiPath, token, nil)
if err != nil {
core.PrintError("Drive unmount", err)
core.ExitWithError(err)
}
defer resp.Body.Close()

respBody, err := io.ReadAll(resp.Body)
if err != nil {
core.PrintError("Drive unmount", fmt.Errorf("failed to read response: %w", err))
core.ExitWithError(err)
}

if resp.StatusCode != http.StatusOK {
handleSandboxAPIError(respBody, resp.StatusCode, "unmount drive")
Expand Down Expand Up @@ -296,18 +294,11 @@ func DriveMountsCmd() *cobra.Command {

sandboxURL, token := resolveSandbox(ctx, sandboxName)

resp, err := sandboxRequest(ctx, http.MethodGet, sandboxURL, "/drives/mount", token, nil)
resp, respBody, err := sandboxRequestWithRetry(ctx, http.MethodGet, sandboxURL, "/drives/mount", token, nil)
if err != nil {
core.PrintError("Drive mounts", err)
core.ExitWithError(err)
}
defer resp.Body.Close()

respBody, err := io.ReadAll(resp.Body)
if err != nil {
core.PrintError("Drive mounts", fmt.Errorf("failed to read response: %w", err))
core.ExitWithError(err)
}

if resp.StatusCode != http.StatusOK {
handleSandboxAPIError(respBody, resp.StatusCode, "list mounted drives")
Expand Down Expand Up @@ -593,6 +584,61 @@ func sandboxRequest(ctx context.Context, method, sandboxURL, path, token string,
return resp, nil
}

// sandboxRequestWithRetry wraps sandboxRequest with exponential backoff retry logic
// for responses marked as retryable by the sandbox API.
func sandboxRequestWithRetry(ctx context.Context, method, sandboxURL, path, token string, body []byte) (*http.Response, []byte, error) {
const (
initialBackoff = 500 * time.Millisecond
maxBackoff = 30 * time.Second
totalBudget = 60 * time.Second
)

deadline := time.Now().Add(totalBudget)
backoff := initialBackoff
retried := false

for {
var bodyReader io.Reader
if body != nil {
bodyReader = bytes.NewReader(body)
}

resp, err := sandboxRequest(ctx, method, sandboxURL, path, token, bodyReader)
if err != nil {
return nil, nil, err
}

respBody, err := io.ReadAll(resp.Body)
resp.Body.Close()
if err != nil {
return nil, nil, err
}

if resp.StatusCode == http.StatusOK {
return resp, respBody, nil
}

// Check if error is retryable
var errResp sandboxErrorResponse
if err := json.Unmarshal(respBody, &errResp); err == nil && errResp.Error.Retryable {
if time.Now().Add(backoff).After(deadline) {
// Budget exhausted, return the error
return resp, respBody, nil
}
if !retried {
core.PrintInfo("Waiting for sandbox to become available...")
retried = true
}
time.Sleep(backoff)
backoff = min(backoff*2, maxBackoff)
continue
}

// Not retryable, return immediately
return resp, respBody, nil
}
}

// handleSandboxAPIError extracts an error message from a sandbox API response and exits.
func handleSandboxAPIError(body []byte, statusCode int, operation string) {
var apiErr sandboxAPIError
Expand Down
Loading