add better error handling and retry

Author: theFong

.github/workflows/validation-lambdalabs.yml

@@ -9,7 +9,6 @@ on:
   pull_request:
     paths:
       - 'internal/lambdalabs/**'
-      - 'pkg/v1/**'
     branches: [ main ]
 
 jobs:
@@ -44,7 +43,7 @@ jobs:
         LAMBDALABS_API_KEY: ${{ secrets.LAMBDALABS_API_KEY }}
       run: |
         cd internal/lambdalabs
-        go test -v -short=false -timeout=20m ./...
+        go test -v -short=false -timeout=30m ./...
         
     - name: Upload test results
       uses: actions/upload-artifact@v4

internal/collections/collections.go

@@ -1,5 +1,11 @@
 package collections
 
+import (
+	"fmt"
+
+	"github.com/cenkalti/backoff/v4"
+)
+
 func Flatten[T any](listOfLists [][]T) []T {
 	result := []T{}
 	for _, list := range listOfLists {
@@ -78,3 +84,15 @@ func Filter[T any](list []T, f func(T) bool) []T {
 func Ptr[T any](x T) *T {
 	return &x
 }
+
+func RetryWithDataAndAttemptCount[T any](o backoff.OperationWithData[T], b backoff.BackOff) (T, error) {
+	attemptCount := 0
+	t, err := backoff.RetryWithData(func() (T, error) {
+		attemptCount++
+		return o()
+	}, b)
+	if err != nil {
+		return t, fmt.Errorf("attemptCount %d: %w", attemptCount, err)
+	}
+	return t, nil
+}

internal/lambdalabs/v1/client.go

@@ -5,9 +5,11 @@ import (
 	"crypto/sha256"
 	"fmt"
 	"net/http"
+	"time"
 
 	openapi "github.com/brevdev/cloud/internal/lambdalabs/gen/lambdalabs"
 	v1 "github.com/brevdev/cloud/pkg/v1"
+	"github.com/cenkalti/backoff/v4"
 )
 
 // LambdaLabsCredential implements the CloudCredential interface for Lambda Labs
@@ -118,3 +120,10 @@ func (c *LambdaLabsClient) makeAuthContext(ctx context.Context) context.Context
 		UserName: c.apiKey,
 	})
 }
+
+func getBackoff() backoff.BackOff {
+	bo := backoff.NewExponentialBackOff()
+	bo.InitialInterval = 1000 * time.Millisecond
+	bo.MaxElapsedTime = 120 * time.Second
+	return bo
+}

internal/lambdalabs/v1/errors.go

@@ -0,0 +1,51 @@
+package v1
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"net/http"
+	"strings"
+
+	openapi "github.com/brevdev/cloud/internal/lambdalabs/gen/lambdalabs"
+	v1 "github.com/brevdev/cloud/pkg/v1"
+	"github.com/cenkalti/backoff/v4"
+)
+
+func handleAPIError(ctx context.Context, resp *http.Response, err error) error {
+	body := ""
+	e, ok := err.(openapi.GenericOpenAPIError)
+	if ok {
+		body = string(e.Body())
+	}
+	if body == "" {
+		bodyBytes, errr := io.ReadAll(resp.Body)
+		if errr != nil {
+			fmt.Printf("Error reading response body: %v\n", errr)
+		}
+		body = string(bodyBytes)
+	}
+	outErr := fmt.Errorf("LambdaLabs API error\n%s\n%s:\nErr: %s\n%s", resp.Request.URL, resp.Status, err.Error(), body)
+	if strings.Contains(body, "instance does not exist") { //nolint:gocritic // ignore
+		return backoff.Permanent(v1.ErrInstanceNotFound)
+	} else if strings.Contains(body, "banned you temporarily") {
+		return outErr
+	} else if resp.StatusCode < 500 && resp.StatusCode != 429 { // 429 Too Many Requests (use back off)
+		return backoff.Permanent(outErr)
+	} else {
+		return outErr
+	}
+}
+
+func handleErrToCloudErr(e error) error {
+	if e == nil {
+		return nil
+	}
+	if strings.Contains(e.Error(), "Not enough capacity") || strings.Contains(e.Error(), "insufficient-capacity") { //nolint:gocritic // ignore
+		return v1.ErrInsufficientResources
+	} else if strings.Contains(e.Error(), "global/invalid-parameters") && strings.Contains(e.Error(), "Region") && strings.Contains(e.Error(), "does not exist") {
+		return v1.ErrInsufficientResources
+	} else {
+		return e
+	}
+}

internal/lambdalabs/v1/instance.go

@@ -66,7 +66,7 @@ func (c *LambdaLabsClient) CreateInstance(ctx context.Context, attrs v1.CreateIn
 		defer func() { _ = httpResp.Body.Close() }()
 	}
 	if err != nil {
-		return nil, fmt.Errorf("failed to launch instance: %w", err)
+		return nil, fmt.Errorf("failed to launch instance: %w", handleErrToCloudErr(err))
 	}
 
 	if len(resp.Data.InstanceIds) != 1 {

internal/lambdalabs/v1/instancetype.go

@@ -90,15 +90,21 @@ func (c *LambdaLabsClient) GetInstanceTypes(ctx context.Context, args v1.GetInst
 }
 
 func (c *LambdaLabsClient) getInstanceTypes(ctx context.Context) (*openapi.InstanceTypes200Response, error) {
-	resp, httpResp, err := c.client.DefaultAPI.InstanceTypes(c.makeAuthContext(ctx)).Execute()
-	if httpResp != nil {
-		defer func() { _ = httpResp.Body.Close() }()
-	}
+	ilr, err := collections.RetryWithDataAndAttemptCount(func() (*openapi.InstanceTypes200Response, error) {
+		res, resp, err := c.client.DefaultAPI.InstanceTypes(c.makeAuthContext(ctx)).Execute()
+		if resp != nil {
+			defer resp.Body.Close() //nolint:errcheck // ignore because using defer (for some reason HandleErrDefer)
+		}
+		if err != nil {
+			return &openapi.InstanceTypes200Response{}, handleAPIError(ctx, resp, err)
+		}
+		return res, nil
+	}, getBackoff())
 	if err != nil {
-		return nil, fmt.Errorf("failed to get instance types: %w", err)
+		return nil, err
 	}
 
-	return resp, nil
+	return ilr, nil
 }
 
 func parseGPUFromDescription(input string) (v1.GPU, error) {