Skip to content

Commit 632df98

Browse files
jvstmejvstme
committed
Test
1 parent 7bc818e commit 632df98

5 files changed

Lines changed: 216 additions & 34 deletions

File tree

‎runner/internal/shim/api/schemas.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,9 @@ type TaskInfoResponse struct {
2929
TerminationReason string `json:"termination_reason"`
3030
TerminationMessage string `json:"termination_message"`
3131
Ports []shim.PortMapping `json:"ports"`
32+
33+
ImagePullProgress *shim.ImagePullProgress `json:"image_pull_progress"`
34+
3235
// The following fields are for debugging only, server doesn't need them
3336
ContainerName string `json:"container_name"`
3437
ContainerID string `json:"container_id"`

‎runner/internal/shim/docker.go‎

Lines changed: 84 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import (
1616
rt "runtime"
1717
"strconv"
1818
"strings"
19+
"sync"
1920
"time"
2021

2122
"github.com/docker/docker/api/types/container"
@@ -50,6 +51,81 @@ const (
5051
LabelValueTrue = "true"
5152
)
5253

54+
// https://github.com/moby/moby/blob/e77ff99ede5ee5952b3a9227863552ae6e5b6fb1/pkg/jsonmessage/jsonmessage.go#L144
55+
// All fields are optional.
56+
type PullMessage struct {
57+
Id string `json:"id"` // layer id
58+
Status string `json:"status"`
59+
ProgressDetail struct {
60+
Current uint64 `json:"current"` // bytes
61+
Total uint64 `json:"total"` // bytes
62+
} `json:"progressDetail"`
63+
ErrorDetail struct {
64+
Message string `json:"message"`
65+
} `json:"errorDetail"`
66+
}
67+
68+
type PullTracker struct {
69+
mu sync.RWMutex
70+
layers map[string]LayerPullProgress
71+
}
72+
73+
func newPullTracker() *PullTracker {
74+
return &PullTracker{layers: make(map[string]LayerPullProgress)}
75+
}
76+
77+
func (t *PullTracker) Update(msg PullMessage) {
78+
if msg.Id == "" {
79+
return
80+
}
81+
t.mu.Lock()
82+
defer t.mu.Unlock()
83+
layer := t.layers[msg.Id]
84+
switch msg.Status {
85+
case "Pulling fs layer", "Waiting", "Verifying Checksum", "Already exists":
86+
// no bytes to update, just track status
87+
case "Downloading":
88+
layer.DownloadedBytes = msg.ProgressDetail.Current
89+
layer.TotalBytes = msg.ProgressDetail.Total
90+
case "Download complete":
91+
layer.DownloadedBytes = layer.TotalBytes
92+
case "Extracting":
93+
layer.ExtractedBytes = msg.ProgressDetail.Current
94+
layer.DownloadedBytes = msg.ProgressDetail.Total
95+
layer.TotalBytes = msg.ProgressDetail.Total
96+
case "Pull complete":
97+
layer.ExtractedBytes = layer.TotalBytes
98+
layer.DownloadedBytes = layer.TotalBytes
99+
default:
100+
// Non-layer events, such as {"status":"Pulling from library/python","id":"3.11"}
101+
return
102+
}
103+
layer.Status = msg.Status
104+
t.layers[msg.Id] = layer
105+
}
106+
107+
func (t *PullTracker) Progress() *ImagePullProgress {
108+
t.mu.RLock()
109+
defer t.mu.RUnlock()
110+
if len(t.layers) == 0 {
111+
return nil
112+
}
113+
p := ImagePullProgress{
114+
IsTotalBytesFinal: true,
115+
Layers: make(map[string]LayerPullProgress, len(t.layers)),
116+
}
117+
for id, l := range t.layers {
118+
if l.TotalBytes == 0 && l.Status != "Already exists" && l.Status != "Pull complete" {
119+
p.IsTotalBytesFinal = false
120+
}
121+
p.DownloadedBytes += l.DownloadedBytes
122+
p.ExtractedBytes += l.ExtractedBytes
123+
p.TotalBytes += l.TotalBytes
124+
p.Layers[id] = l
125+
}
126+
return &p
127+
}
128+
53129
type DockerRunner struct {
54130
client *docker.Client
55131
dockerParams DockerParameters
@@ -239,6 +315,7 @@ func (d *DockerRunner) TaskInfo(taskID string) TaskInfo {
239315
ContainerName: task.containerName,
240316
ContainerID: task.containerID,
241317
GpuIDs: task.gpuIDs,
318+
ImagePullProgress: task.pullTracker.Progress(),
242319
}
243320
}
244321

@@ -350,7 +427,7 @@ func (d *DockerRunner) Run(ctx context.Context, taskID string) error {
350427
// Although it's called "runner dir", we also use it for shim task-related data.
351428
// Maybe we should rename it to "task dir" (including the `/root/.dstack/runners` dir on the host).
352429
pullLogPath := filepath.Join(runnerDir, "pull.log")
353-
if err = pullImage(pullCtx, d.client, cfg, pullLogPath); err != nil {
430+
if err = pullImage(pullCtx, d.client, cfg, pullLogPath, task.pullTracker); err != nil {
354431
errMessage := fmt.Sprintf("pullImage error: %s", err.Error())
355432
log.Error(ctx, errMessage)
356433
task.SetStatusTerminated(string(types.TerminationReasonCreatingContainerError), errMessage)
@@ -670,7 +747,7 @@ func mountDisk(ctx context.Context, deviceName, mountPoint string, fsRootPerms o
670747
return nil
671748
}
672749

673-
func pullImage(ctx context.Context, client docker.APIClient, taskConfig TaskConfig, logPath string) error {
750+
func pullImage(ctx context.Context, client docker.APIClient, taskConfig TaskConfig, logPath string, tracker *PullTracker) error {
674751
if !strings.Contains(taskConfig.ImageName, ":") {
675752
taskConfig.ImageName += ":latest"
676753
}
@@ -710,26 +787,8 @@ func pullImage(ctx context.Context, client docker.APIClient, taskConfig TaskConf
710787

711788
teeReader := io.TeeReader(reader, logFile)
712789

713-
current := make(map[string]uint)
714-
total := make(map[string]uint)
715-
716790
// dockerd reports pulling progress as a stream of JSON Lines. The format of records is not documented in the API documentation,
717791
// although it's occasionally mentioned, e.g., https://docs.docker.com/reference/api/engine/version-history/#v148-api-changes
718-
719-
// https://github.com/moby/moby/blob/e77ff99ede5ee5952b3a9227863552ae6e5b6fb1/pkg/jsonmessage/jsonmessage.go#L144
720-
// All fields are optional
721-
type PullMessage struct {
722-
Id string `json:"id"` // layer id
723-
Status string `json:"status"`
724-
ProgressDetail struct {
725-
Current uint `json:"current"` // bytes
726-
Total uint `json:"total"` // bytes
727-
} `json:"progressDetail"`
728-
ErrorDetail struct {
729-
Message string `json:"message"`
730-
} `json:"errorDetail"`
731-
}
732-
733792
var pullCompleted bool
734793
pullErrors := make([]string, 0)
735794

@@ -740,13 +799,7 @@ func pullImage(ctx context.Context, client docker.APIClient, taskConfig TaskConf
740799
if err := json.Unmarshal(line, &pullMessage); err != nil {
741800
continue
742801
}
743-
if pullMessage.Status == "Downloading" {
744-
current[pullMessage.Id] = pullMessage.ProgressDetail.Current
745-
total[pullMessage.Id] = pullMessage.ProgressDetail.Total
746-
}
747-
if pullMessage.Status == "Download complete" {
748-
current[pullMessage.Id] = total[pullMessage.Id]
749-
}
802+
tracker.Update(pullMessage)
750803
if pullMessage.ErrorDetail.Message != "" {
751804
log.Error(ctx, "error pulling image", "name", taskConfig.ImageName, "err", pullMessage.ErrorDetail.Message)
752805
pullErrors = append(pullErrors, pullMessage.ErrorDetail.Message)
@@ -764,13 +817,10 @@ func pullImage(ctx context.Context, client docker.APIClient, taskConfig TaskConf
764817
}
765818

766819
duration := time.Since(startTime)
767-
var currentBytes uint
768-
var totalBytes uint
769-
for _, v := range current {
770-
currentBytes += v
771-
}
772-
for _, v := range total {
773-
totalBytes += v
820+
p := tracker.Progress()
821+
var currentBytes, totalBytes uint64
822+
if p != nil {
823+
currentBytes, totalBytes = p.DownloadedBytes, p.TotalBytes
774824
}
775825
speed := bytesize.New(float64(currentBytes) / duration.Seconds())
776826

‎runner/internal/shim/docker_test.go‎

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -155,3 +155,112 @@ func createTaskConfig(t *testing.T) TaskConfig {
155155
ImageName: "ubuntu",
156156
}
157157
}
158+
159+
func pullMsg(id, status string, current, total uint64) PullMessage {
160+
m := PullMessage{Id: id, Status: status}
161+
m.ProgressDetail.Current = current
162+
m.ProgressDetail.Total = total
163+
return m
164+
}
165+
166+
func TestPullTracker_Empty(t *testing.T) {
167+
tracker := newPullTracker()
168+
assert.Nil(t, tracker.Progress())
169+
}
170+
171+
func TestPullTracker_AlreadyExists(t *testing.T) {
172+
tracker := newPullTracker()
173+
tracker.Update(PullMessage{Id: "3.11", Status: "Pulling from library/python"})
174+
for _, id := range []string{"aaa", "bbb", "ccc"} {
175+
tracker.Update(PullMessage{Id: id, Status: "Already exists"})
176+
}
177+
tracker.Update(PullMessage{Status: "Digest: sha256:***"})
178+
tracker.Update(PullMessage{Status: "Status: Image is up to date for python:3.11"})
179+
p := tracker.Progress()
180+
require.NotNil(t, p)
181+
assert.Equal(t, uint64(0), p.DownloadedBytes)
182+
assert.Equal(t, uint64(0), p.ExtractedBytes)
183+
assert.Equal(t, uint64(0), p.TotalBytes)
184+
assert.True(t, p.IsTotalBytesFinal)
185+
}
186+
187+
func TestPullTracker_FullPull(t *testing.T) {
188+
const sizeA, sizeB uint64 = 111, 222
189+
190+
tracker := newPullTracker()
191+
tracker.Update(PullMessage{Id: "3.11", Status: "Pulling from library/python"})
192+
tracker.Update(PullMessage{Id: "aaa", Status: "Pulling fs layer"})
193+
tracker.Update(PullMessage{Id: "bbb", Status: "Pulling fs layer"})
194+
tracker.Update(PullMessage{Id: "aaa", Status: "Waiting"})
195+
tracker.Update(PullMessage{Id: "bbb", Status: "Waiting"})
196+
197+
// Layers announced but sizes unknown yet
198+
p := tracker.Progress()
199+
require.NotNil(t, p)
200+
assert.Equal(t, uint64(0), p.DownloadedBytes)
201+
assert.Equal(t, uint64(0), p.ExtractedBytes)
202+
assert.Equal(t, uint64(0), p.TotalBytes)
203+
assert.False(t, p.IsTotalBytesFinal)
204+
205+
// Both layers start downloading - sizes now known
206+
tracker.Update(pullMsg("aaa", "Downloading", 100, sizeA))
207+
tracker.Update(pullMsg("bbb", "Downloading", 200, sizeB))
208+
209+
p = tracker.Progress()
210+
assert.Equal(t, uint64(300), p.DownloadedBytes)
211+
assert.Equal(t, uint64(0), p.ExtractedBytes)
212+
assert.Equal(t, sizeA+sizeB, p.TotalBytes)
213+
assert.True(t, p.IsTotalBytesFinal)
214+
215+
// Downloads complete
216+
tracker.Update(pullMsg("aaa", "Downloading", sizeA, sizeA))
217+
tracker.Update(PullMessage{Id: "aaa", Status: "Download complete"})
218+
tracker.Update(pullMsg("bbb", "Downloading", sizeB, sizeB))
219+
tracker.Update(PullMessage{Id: "bbb", Status: "Download complete"})
220+
221+
p = tracker.Progress()
222+
assert.Equal(t, sizeA+sizeB, p.DownloadedBytes)
223+
assert.Equal(t, uint64(0), p.ExtractedBytes)
224+
assert.Equal(t, sizeA+sizeB, p.TotalBytes)
225+
assert.True(t, p.IsTotalBytesFinal)
226+
227+
// Both layers start extracting
228+
tracker.Update(pullMsg("aaa", "Extracting", 100, sizeA))
229+
tracker.Update(pullMsg("bbb", "Extracting", 200, sizeB))
230+
231+
p = tracker.Progress()
232+
assert.Equal(t, sizeA+sizeB, p.DownloadedBytes)
233+
assert.Equal(t, uint64(300), p.ExtractedBytes)
234+
assert.Equal(t, sizeA+sizeB, p.TotalBytes)
235+
assert.True(t, p.IsTotalBytesFinal)
236+
237+
// Extractions complete
238+
tracker.Update(pullMsg("aaa", "Extracting", sizeA, sizeA))
239+
tracker.Update(PullMessage{Id: "aaa", Status: "Pull complete"})
240+
tracker.Update(pullMsg("bbb", "Extracting", sizeB, sizeB))
241+
tracker.Update(PullMessage{Id: "bbb", Status: "Pull complete"})
242+
tracker.Update(PullMessage{Status: "Digest: sha256:***"})
243+
tracker.Update(PullMessage{Status: "Status: Downloaded newer image for python:3.11"})
244+
245+
p = tracker.Progress()
246+
assert.Equal(t, sizeA+sizeB, p.DownloadedBytes)
247+
assert.Equal(t, sizeA+sizeB, p.ExtractedBytes)
248+
assert.Equal(t, sizeA+sizeB, p.TotalBytes)
249+
assert.True(t, p.IsTotalBytesFinal)
250+
}
251+
252+
func TestPullTracker_MixedLayerStatuses(t *testing.T) {
253+
tracker := newPullTracker()
254+
255+
tracker.Update(PullMessage{Id: "layer-exists", Status: "Already exists"})
256+
tracker.Update(pullMsg("layer-downloading", "Downloading", 50, 100))
257+
tracker.Update(pullMsg("layer-extracting", "Extracting", 100, 200))
258+
tracker.Update(PullMessage{Id: "layer-waiting", Status: "Waiting"})
259+
260+
p := tracker.Progress()
261+
require.NotNil(t, p)
262+
assert.Equal(t, uint64(50+200), p.DownloadedBytes)
263+
assert.Equal(t, uint64(100), p.ExtractedBytes)
264+
assert.Equal(t, uint64(100+200), p.TotalBytes)
265+
assert.False(t, p.IsTotalBytesFinal) // layer-waiting size unknown
266+
}

‎runner/internal/shim/models.go‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,12 +108,28 @@ type TaskListItem struct {
108108
Status TaskStatus `json:"status"`
109109
}
110110

111+
type LayerPullProgress struct {
112+
Status string `json:"status"`
113+
DownloadedBytes uint64 `json:"downloaded_bytes"`
114+
ExtractedBytes uint64 `json:"extracted_bytes"`
115+
TotalBytes uint64 `json:"total_bytes"`
116+
}
117+
118+
type ImagePullProgress struct {
119+
DownloadedBytes uint64 `json:"downloaded_bytes"`
120+
ExtractedBytes uint64 `json:"extracted_bytes"`
121+
TotalBytes uint64 `json:"total_bytes"`
122+
IsTotalBytesFinal bool `json:"is_total_bytes_final"`
123+
Layers map[string]LayerPullProgress `json:"layers"`
124+
}
125+
111126
type TaskInfo struct {
112127
ID string
113128
Status TaskStatus
114129
TerminationReason string
115130
TerminationMessage string
116131
Ports []PortMapping
132+
ImagePullProgress *ImagePullProgress
117133
ContainerName string
118134
ContainerID string
119135
GpuIDs []string

‎runner/internal/shim/task.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,8 @@ type Task struct {
4242
ports []PortMapping
4343
runnerDir string // path on host mapped to consts.RunnerDir in container
4444

45+
pullTracker *PullTracker
46+
4547
mu *sync.Mutex
4648
}
4749

@@ -128,6 +130,7 @@ func NewTask(id string, status TaskStatus, containerName string, containerID str
128130
runnerDir: runnerDir,
129131
gpuIDs: gpuIDs,
130132
ports: ports,
133+
pullTracker: newPullTracker(),
131134
mu: &sync.Mutex{},
132135
}
133136
}
@@ -138,6 +141,7 @@ func NewTaskFromConfig(cfg TaskConfig) Task {
138141
Status: TaskStatusPending,
139142
config: cfg,
140143
containerName: generateUniqueName(cfg.Name, cfg.ID),
144+
pullTracker: newPullTracker(),
141145
mu: &sync.Mutex{},
142146
}
143147
}

0 commit comments

Comments
 (0)