Skip to content

Commit ac1812a

Browse files
Edwinhr716Edwinhr716
committed
prototype error handling
1 parent 97d2c49 commit ac1812a

10 files changed

Lines changed: 113 additions & 25 deletions

File tree

‎pkg/accelerator-orchestrator/api/v1alpha1/accelerator_orchestrator.pb.go‎

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎pkg/accelerator-orchestrator/api/v1alpha1/accelerator_orchestrator_grpc.pb.go‎

Lines changed: 6 additions & 6 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎pkg/snapshot-agent/api/v1alpha1/snapshot_agent.pb.go‎

Lines changed: 9 additions & 6 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎pkg/snapshot-agent/api/v1alpha1/snapshot_agent.proto‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,7 @@ enum JobState {
105105
JOB_STATE_TRANSITIONING = 3;
106106
JOB_STATE_SAVED = 4;
107107
JOB_STATE_FAULTED = 5;
108+
JOB_STATE_NOT_FOUND = 6;
108109
}
109110

110111
message JobStatus {

‎pkg/snapshot-agent/api/v1alpha1/snapshot_agent_grpc.pb.go‎

Lines changed: 6 additions & 7 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎pkg/snapshot-agent/backends/cuda-checkpoint.go‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,15 +2,21 @@ package backends
22

33
import (
44
"context"
5+
"errors"
56
"fmt"
67
"log/slog"
78
"os/exec"
9+
"strings"
810
"sync"
911
"time"
1012

1113
"github.com/NVIDIA/go-nvml/pkg/nvml"
1214
)
1315

16+
// ErrPIDNotFound is returned when cuda-checkpoint fails because the target PID was not found.
17+
var ErrPIDNotFound = errors.New("pid not found")
18+
19+
1420
type nvmlClient interface {
1521
Init() nvml.Return
1622
Shutdown() nvml.Return
@@ -95,7 +101,11 @@ func (c *CudaCheckpoint) getCudaCheckpointPath() string {
95101

96102
func (c *CudaCheckpoint) runSudoCommand(ctx context.Context, name string, args ...string) error {
97103
if out, err := c.execCommand(ctx, name, args...); err != nil {
98-
return fmt.Errorf("command failed: %w, output: %s", err, string(out))
104+
outStr := string(out)
105+
if strings.Contains(outStr, "No such process") || strings.Contains(outStr, "process not found") || strings.Contains(outStr, "invalid PID") {
106+
return fmt.Errorf("command failed: %w: %s", ErrPIDNotFound, outStr)
107+
}
108+
return fmt.Errorf("command failed: %w, output: %s", err, outStr)
99109
}
100110
return nil
101111
}

‎pkg/snapshot-agent/backends/cuda_checkpoint_test.go‎

Lines changed: 45 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package backends_test
22

33
import (
44
"context"
5+
"errors"
56
"fmt"
67
"testing"
78

@@ -32,7 +33,9 @@ func TestSnapshot(t *testing.T) {
3233
name string
3334
pids []string
3435
execErr error
36+
execOut []byte
3537
expectedErr bool
38+
errIs error
3639
}{
3740
{
3841
name: "Success",
@@ -46,19 +49,38 @@ func TestSnapshot(t *testing.T) {
4649
execErr: fmt.Errorf("exec error"),
4750
expectedErr: true,
4851
},
52+
{
53+
name: "PIDNotFound_NoSuchProcess",
54+
pids: []string{"123"},
55+
execErr: fmt.Errorf("exec error"),
56+
execOut: []byte("Unable to detach from 123: No such process"),
57+
expectedErr: true,
58+
errIs: backends.ErrPIDNotFound,
59+
},
60+
{
61+
name: "PIDNotFound_ProcessNotFound",
62+
pids: []string{"123"},
63+
execErr: fmt.Errorf("exec error"),
64+
execOut: []byte("process not found"),
65+
expectedErr: true,
66+
errIs: backends.ErrPIDNotFound,
67+
},
4968
}
5069

5170
for _, tt := range tests {
5271
t.Run(tt.name, func(t *testing.T) {
5372
c := backends.NewCudaCheckpoint()
5473
c.SetExecCommand(func(ctx context.Context, name string, args ...string) ([]byte, error) {
55-
return nil, tt.execErr
74+
return tt.execOut, tt.execErr
5675
})
5776

5877
err := c.Snapshot(context.Background(), tt.pids)
5978
if (err != nil) != tt.expectedErr {
6079
t.Errorf("Snapshot() error = %v, expectedErr %v", err, tt.expectedErr)
6180
}
81+
if tt.errIs != nil && !errors.Is(err, tt.errIs) {
82+
t.Errorf("Snapshot() error = %v, expected to be %v", err, tt.errIs)
83+
}
6284
})
6385
}
6486
}
@@ -68,7 +90,9 @@ func TestRestore(t *testing.T) {
6890
name string
6991
pids []string
7092
execErr error
93+
execOut []byte
7194
expectedErr bool
95+
errIs error
7296
}{
7397
{
7498
name: "Success",
@@ -88,19 +112,38 @@ func TestRestore(t *testing.T) {
88112
execErr: fmt.Errorf("exec error"),
89113
expectedErr: true,
90114
},
115+
{
116+
name: "PIDNotFound_NoSuchProcess",
117+
pids: []string{"123"},
118+
execErr: fmt.Errorf("exec error"),
119+
execOut: []byte("Unable to detach from 123: No such process"),
120+
expectedErr: true,
121+
errIs: backends.ErrPIDNotFound,
122+
},
123+
{
124+
name: "PIDNotFound_ProcessNotFound",
125+
pids: []string{"123"},
126+
execErr: fmt.Errorf("exec error"),
127+
execOut: []byte("process not found"),
128+
expectedErr: true,
129+
errIs: backends.ErrPIDNotFound,
130+
},
91131
}
92132

93133
for _, tt := range tests {
94134
t.Run(tt.name, func(t *testing.T) {
95135
c := backends.NewCudaCheckpoint()
96136
c.SetExecCommand(func(ctx context.Context, name string, args ...string) ([]byte, error) {
97-
return nil, tt.execErr
137+
return tt.execOut, tt.execErr
98138
})
99139

100140
err := c.Restore(context.Background(), tt.pids)
101141
if (err != nil) != tt.expectedErr {
102142
t.Errorf("Restore() error = %v, expectedErr %v", err, tt.expectedErr)
103143
}
144+
if tt.errIs != nil && !errors.Is(err, tt.errIs) {
145+
t.Errorf("Restore() error = %v, expected to be %v", err, tt.errIs)
146+
}
104147
})
105148
}
106149
}

‎pkg/snapshot-agent/server/server.go‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package server
22

33
import (
44
"context"
5+
"errors"
56
"fmt"
67
"log/slog"
78
"net"
@@ -95,6 +96,9 @@ func (s *Server) Snapshot(ctx context.Context, req *pb.SnapshotRequest) (*pb.Sna
9596

9697
err = backend.Snapshot(bgCtx, allPIDStrings)
9798
if err != nil {
99+
if errors.Is(err, backends.ErrPIDNotFound) {
100+
return fmt.Errorf("failed to snapshot job %s: %w", req.GetJobId(), sm.ErrJobNotFound)
101+
}
98102
return fmt.Errorf("failed to snapshot job %s: %w", req.GetJobId(), err)
99103
}
100104

@@ -151,6 +155,9 @@ func (s *Server) Restore(ctx context.Context, req *pb.RestoreRequest) (*pb.Resto
151155

152156
slog.InfoContext(bgCtx, "Restoring PIDs", "pids", pidStrings, "backend", backendType)
153157
if err := backend.Restore(bgCtx, pidStrings); err != nil {
158+
if errors.Is(err, backends.ErrPIDNotFound) {
159+
return fmt.Errorf("failed to restore job %s: %w", req.GetJobId(), sm.ErrJobNotFound)
160+
}
154161
return fmt.Errorf("failed to restore job %s: %w", req.GetJobId(), err)
155162
}
156163
return nil

‎pkg/snapshot-agent/state-machine/state-manager.go‎

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package statemachine
22

33
import (
4+
"errors"
45
"sync"
56
"time"
67

@@ -10,6 +11,9 @@ import (
1011
"google.golang.org/grpc/status"
1112
)
1213

14+
// ErrJobNotFound is returned when the job's processes cannot be found.
15+
var ErrJobNotFound = errors.New("job not found")
16+
1317
// OpType represents the type of operation (Snapshot or Restore).
1418
type OpType string
1519

@@ -126,7 +130,11 @@ func (sm *StateManager) StartSnapshot(jobID, group string, worker func() error)
126130
if err != nil {
127131
op.Status = pb.OperationStatus_OPERATION_STATUS_FAILED
128132
op.Error = err.Error()
129-
job.State = pb.JobState_JOB_STATE_FAULTED
133+
if errors.Is(err, ErrJobNotFound) {
134+
job.State = pb.JobState_JOB_STATE_NOT_FOUND
135+
} else {
136+
job.State = pb.JobState_JOB_STATE_FAULTED
137+
}
130138
} else {
131139
op.Status = pb.OperationStatus_OPERATION_STATUS_COMPLETE
132140
op.StorageBytes = 1024
@@ -191,7 +199,11 @@ func (sm *StateManager) StartRestore(jobID, group string, worker func() error) (
191199
if err != nil {
192200
op.Status = pb.OperationStatus_OPERATION_STATUS_FAILED
193201
op.Error = err.Error()
194-
job.State = pb.JobState_JOB_STATE_FAULTED
202+
if errors.Is(err, ErrJobNotFound) {
203+
job.State = pb.JobState_JOB_STATE_NOT_FOUND
204+
} else {
205+
job.State = pb.JobState_JOB_STATE_FAULTED
206+
}
195207
} else {
196208
op.Status = pb.OperationStatus_OPERATION_STATUS_COMPLETE
197209
job.State = pb.JobState_JOB_STATE_RUNNING

0 commit comments

Comments
 (0)