Test

Author: jvstme

runner/internal/shim/api/schemas.go

@@ -29,6 +29,10 @@ type TaskInfoResponse struct {
 	TerminationReason  string             `json:"termination_reason"`
 	TerminationMessage string             `json:"termination_message"`
 	Ports              []shim.PortMapping `json:"ports"`
+
+	ImagePullCurrentBytes *int64 `json:"image_pull_current_bytes"`
+	ImagePullTotalBytes   *int64 `json:"image_pull_total_bytes"`
+
 	// The following fields are for debugging only, server doesn't need them
 	ContainerName string   `json:"container_name"`
 	ContainerID   string   `json:"container_id"`

runner/internal/shim/docker.go

@@ -231,14 +231,16 @@ func (d *DockerRunner) TaskInfo(taskID string) TaskInfo {
 		return TaskInfo{}
 	}
 	return TaskInfo{
-		ID:                 task.ID,
-		Status:             task.Status,
-		TerminationReason:  task.TerminationReason,
-		TerminationMessage: task.TerminationMessage,
-		Ports:              task.ports,
-		ContainerName:      task.containerName,
-		ContainerID:        task.containerID,
-		GpuIDs:             task.gpuIDs,
+		ID:                    task.ID,
+		Status:                task.Status,
+		TerminationReason:     task.TerminationReason,
+		TerminationMessage:    task.TerminationMessage,
+		Ports:                 task.ports,
+		ContainerName:         task.containerName,
+		ContainerID:           task.containerID,
+		GpuIDs:                task.gpuIDs,
+		ImagePullCurrentBytes: task.imagePullCurrentBytes,
+		ImagePullTotalBytes:   task.imagePullTotalBytes,
 	}
 }
 
@@ -350,7 +352,14 @@ func (d *DockerRunner) Run(ctx context.Context, taskID string) error {
 	// Although it's called "runner dir", we also use it for shim task-related data.
 	// Maybe we should rename it to "task dir" (including the `/root/.dstack/runners` dir on the host).
 	pullLogPath := filepath.Join(runnerDir, "pull.log")
-	if err = pullImage(pullCtx, d.client, cfg, pullLogPath); err != nil {
+	onPullProgress := func(currentBytes, totalBytes int64) {
+		task.imagePullCurrentBytes = &currentBytes
+		task.imagePullTotalBytes = &totalBytes
+		if updateErr := d.tasks.Update(task); updateErr != nil {
+			log.Debug(ctx, "pull progress update skipped", "err", updateErr)
+		}
+	}
+	if err = pullImage(pullCtx, d.client, cfg, pullLogPath, onPullProgress); err != nil {
 		errMessage := fmt.Sprintf("pullImage error: %s", err.Error())
 		log.Error(ctx, errMessage)
 		task.SetStatusTerminated(string(types.TerminationReasonCreatingContainerError), errMessage)
@@ -670,7 +679,7 @@ func mountDisk(ctx context.Context, deviceName, mountPoint string, fsRootPerms o
 	return nil
 }
 
-func pullImage(ctx context.Context, client docker.APIClient, taskConfig TaskConfig, logPath string) error {
+func pullImage(ctx context.Context, client docker.APIClient, taskConfig TaskConfig, logPath string, onProgress func(currentBytes, totalBytes int64)) error {
 	if !strings.Contains(taskConfig.ImageName, ":") {
 		taskConfig.ImageName += ":latest"
 	}
@@ -730,6 +739,20 @@ func pullImage(ctx context.Context, client docker.APIClient, taskConfig TaskConf
 		} `json:"errorDetail"`
 	}
 
+	reportProgress := func() {
+		if onProgress == nil {
+			return
+		}
+		var cb, tb int64
+		for _, v := range current {
+			cb += int64(v)
+		}
+		for _, v := range total {
+			tb += int64(v)
+		}
+		onProgress(cb, tb)
+	}
+
 	var pullCompleted bool
 	pullErrors := make([]string, 0)
 
@@ -743,9 +766,11 @@ func pullImage(ctx context.Context, client docker.APIClient, taskConfig TaskConf
 		if pullMessage.Status == "Downloading" {
 			current[pullMessage.Id] = pullMessage.ProgressDetail.Current
 			total[pullMessage.Id] = pullMessage.ProgressDetail.Total
+			reportProgress()
 		}
 		if pullMessage.Status == "Download complete" {
 			current[pullMessage.Id] = total[pullMessage.Id]
+			reportProgress()
 		}
 		if pullMessage.ErrorDetail.Message != "" {
 			log.Error(ctx, "error pulling image", "name", taskConfig.ImageName, "err", pullMessage.ErrorDetail.Message)

runner/internal/shim/models.go

@@ -109,12 +109,14 @@ type TaskListItem struct {
 }
 
 type TaskInfo struct {
-	ID                 string
-	Status             TaskStatus
-	TerminationReason  string
-	TerminationMessage string
-	Ports              []PortMapping
-	ContainerName      string
-	ContainerID        string
-	GpuIDs             []string
+	ID                    string
+	Status                TaskStatus
+	TerminationReason     string
+	TerminationMessage    string
+	Ports                 []PortMapping
+	ImagePullCurrentBytes *int64
+	ImagePullTotalBytes   *int64
+	ContainerName         string
+	ContainerID           string
+	GpuIDs                []string
 }

runner/internal/shim/task.go

@@ -42,6 +42,9 @@ type Task struct {
 	ports         []PortMapping
 	runnerDir     string // path on host mapped to consts.RunnerDir in container
 
+	imagePullCurrentBytes *int64
+	imagePullTotalBytes   *int64
+
 	mu *sync.Mutex
 }
 
@@ -75,7 +78,8 @@ func (t *Task) IsTransitionAllowed(toStatus TaskStatus) bool {
 	case TaskStatusPreparing:
 		return t.Status == TaskStatusPending
 	case TaskStatusPulling:
-		return t.Status == TaskStatusPreparing
+		// allow pulling->pulling to update pull progress
+		return t.Status == TaskStatusPulling || t.Status == TaskStatusPreparing
 	case TaskStatusCreating:
 		return t.Status == TaskStatusPulling
 	case TaskStatusRunning:

runner/internal/shim/task_test.go

@@ -93,6 +93,7 @@ func TestTask_IsTransitionAllowed_true(t *testing.T) {
 		{TaskStatusPending, TaskStatusTerminated},
 		{TaskStatusPreparing, TaskStatusPulling},
 		{TaskStatusPreparing, TaskStatusTerminated},
+		{TaskStatusPulling, TaskStatusPulling},
 		{TaskStatusPulling, TaskStatusCreating},
 		{TaskStatusPulling, TaskStatusTerminated},
 		{TaskStatusCreating, TaskStatusRunning},