Skip to content

Commit 641739c

Browse files
elwinchengelwincheng
committed
API for docker container and volumes
1 parent c00a898 commit 641739c

File tree

2 files changed

+350
-0
lines changed

2 files changed

+350
-0
lines changed

‎src/images/serve_image.go‎

Lines changed: 165 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,165 @@
1+
package main
2+
3+
import (
4+
"context"
5+
"fmt"
6+
"io"
7+
"os"
8+
9+
"github.com/docker/docker/api/types/container"
10+
"github.com/docker/docker/api/types/mount"
11+
"github.com/docker/docker/api/types/volume"
12+
"github.com/docker/docker/client"
13+
)
14+
15+
type ContainerMgr struct {
16+
ctx context.Context
17+
cli *client.Client
18+
containerLimit int
19+
volumeLimit int
20+
containers map[string]struct{}
21+
volumes map[string]struct{}
22+
}
23+
24+
func NewContainerMgr(client *client.Client, containerLimit, volumeLimit int) *ContainerMgr {
25+
return &ContainerMgr{
26+
ctx: context.Background(),
27+
cli: client,
28+
containerLimit: containerLimit,
29+
volumeLimit: volumeLimit,
30+
containers: make(map[string]struct{}),
31+
volumes: make(map[string]struct{}),
32+
}
33+
}
34+
35+
func (mgr *ContainerMgr) stopContainer(containerID string) {
36+
ctx := mgr.ctx
37+
cli := mgr.cli
38+
39+
err := cli.ContainerStop(ctx, containerID, container.StopOptions{})
40+
if err != nil {
41+
panic(err)
42+
}
43+
44+
}
45+
46+
func (mgr *ContainerMgr) removeContainer(containerID string) error {
47+
ctx := mgr.ctx
48+
cli := mgr.cli
49+
err := cli.ContainerRemove(ctx, containerID, container.RemoveOptions{RemoveVolumes: true})
50+
if err != nil {
51+
return err
52+
}
53+
delete(mgr.containers, containerID)
54+
return nil
55+
}
56+
57+
func (mgr *ContainerMgr) createVolume(volumeName string) (volume.Volume, error) {
58+
if len(mgr.volumes) >= mgr.volumeLimit {
59+
return volume.Volume{}, fmt.Errorf("volume limit reached")
60+
}
61+
ctx := mgr.ctx
62+
cli := mgr.cli
63+
64+
vol, err := cli.VolumeCreate(ctx, volume.CreateOptions{
65+
Name: volumeName, // You can leave this empty for a random name
66+
})
67+
if err != nil {
68+
return volume.Volume{}, err
69+
}
70+
mgr.volumes[vol.Name] = struct{}{}
71+
return vol, nil
72+
}
73+
74+
func (mgr *ContainerMgr) removeVolume(volumeName string, force bool) error {
75+
ctx := mgr.ctx
76+
cli := mgr.cli
77+
78+
vols, _ := cli.VolumeList(ctx, volume.ListOptions{})
79+
found := false
80+
for _, v := range vols.Volumes {
81+
if v.Name == volumeName {
82+
found = true
83+
break
84+
}
85+
}
86+
if !found {
87+
return fmt.Errorf("volume %s does not exist", volumeName)
88+
}
89+
90+
err := cli.VolumeRemove(ctx, volumeName, force)
91+
if err != nil {
92+
return err
93+
}
94+
delete(mgr.volumes, volumeName)
95+
return nil
96+
}
97+
98+
func (mgr *ContainerMgr) runContainerCuda(volumeName string) (string, error) {
99+
if len(mgr.containers) >= mgr.containerLimit {
100+
return "", fmt.Errorf("container limit reached")
101+
}
102+
ctx := mgr.ctx
103+
cli := mgr.cli
104+
105+
resp, err := cli.ContainerCreate(ctx, &container.Config{
106+
Image: "pytorch-cuda",
107+
Cmd: []string{"sleep", "1000"},
108+
}, &container.HostConfig{
109+
Runtime: "nvidia",
110+
Mounts: []mount.Mount{
111+
{
112+
Type: mount.TypeVolume,
113+
Source: volumeName,
114+
Target: "/data",
115+
},
116+
},
117+
}, nil, nil, "")
118+
vols, _ := cli.VolumeList(ctx, volume.ListOptions{})
119+
found := false
120+
for _, v := range vols.Volumes {
121+
if v.Name == volumeName {
122+
found = true
123+
break
124+
}
125+
}
126+
if !found {
127+
return "", fmt.Errorf("volume %s does not exist", volumeName)
128+
}
129+
if err != nil {
130+
return "", err
131+
}
132+
mgr.containers[resp.ID] = struct{}{}
133+
if err := cli.ContainerStart(ctx, resp.ID, container.StartOptions{}); err != nil {
134+
return "", err
135+
}
136+
137+
out, err := cli.ContainerLogs(ctx, resp.ID, container.LogsOptions{ShowStdout: true})
138+
if err != nil {
139+
panic(err)
140+
}
141+
142+
io.Copy(os.Stdout, out)
143+
return resp.ID, nil
144+
}
145+
146+
func main() {
147+
cli, err := client.NewClientWithOpts(client.FromEnv)
148+
if err != nil {
149+
panic(err)
150+
}
151+
152+
// Create a Docker volume
153+
containerMgr := NewContainerMgr(cli, 10, 10)
154+
volumeName := "my_volume1"
155+
156+
containerMgr.createVolume(volumeName)
157+
id, err := containerMgr.runContainerCuda(volumeName)
158+
if err != nil {
159+
fmt.Errorf("Failed to start container: %v", err.Error())
160+
}
161+
containerMgr.stopContainer(id)
162+
containerMgr.removeContainer(id)
163+
containerMgr.removeVolume(volumeName, true)
164+
165+
}

‎src/images/serve_image_test.go‎

Lines changed: 185 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,185 @@
1+
package main
2+
3+
import (
4+
"fmt"
5+
"testing"
6+
7+
"github.com/docker/docker/api/types/volume"
8+
"github.com/docker/docker/client"
9+
)
10+
11+
func setupMgr(t *testing.T) *ContainerMgr {
12+
cli, err := client.NewClientWithOpts(client.FromEnv)
13+
if err != nil {
14+
t.Fatalf("Failed to create Docker client: %v", err)
15+
}
16+
return NewContainerMgr(cli, 10, 100)
17+
}
18+
19+
// T1: create a volume, check exists, delete, check not exists
20+
func TestCreateDeleteVolume(t *testing.T) {
21+
mgr := setupMgr(t)
22+
volName := "test_volume_t1"
23+
mgr.createVolume(volName)
24+
vols, _ := mgr.cli.VolumeList(mgr.ctx, volume.ListOptions{})
25+
found := false
26+
for _, v := range vols.Volumes {
27+
if v.Name == volName {
28+
found = true
29+
break
30+
}
31+
}
32+
if !found {
33+
t.Errorf("Volume %s not found after creation", volName)
34+
}
35+
mgr.removeVolume(volName, true)
36+
vols, _ = mgr.cli.VolumeList(mgr.ctx, volume.ListOptions{})
37+
for _, v := range vols.Volumes {
38+
if v.Name == volName {
39+
t.Errorf("Volume %s still exists after deletion", volName)
40+
}
41+
}
42+
}
43+
44+
// T3: create a volume with same name twice (should not fail)
45+
func TestCreateVolumeTwice(t *testing.T) {
46+
mgr := setupMgr(t)
47+
volName := "test_volume_t3"
48+
mgr.createVolume(volName)
49+
defer mgr.removeVolume(volName, true)
50+
mgr.createVolume(volName) // Should not fail
51+
}
52+
53+
// T4: remove volume that doesn't exist (should fail or panic)
54+
func TestRemoveNonexistentVolume(t *testing.T) {
55+
mgr := setupMgr(t)
56+
err := mgr.removeVolume("nonexistent_volume_t4", true)
57+
if err == nil {
58+
t.Errorf("Expected error when removing nonexistent volume, but no error")
59+
}
60+
}
61+
62+
// T5: remove volume in use (should fail or panic)
63+
func TestRemoveVolumeInUse(t *testing.T) {
64+
mgr := setupMgr(t)
65+
volName := "test_volume_t5"
66+
mgr.createVolume(volName)
67+
containerID, err := mgr.runContainerCuda(volName)
68+
if err != nil {
69+
fmt.Errorf("Failed to start container: %v", err.Error())
70+
}
71+
defer func() {
72+
mgr.stopContainer(containerID)
73+
mgr.removeContainer(containerID)
74+
mgr.removeVolume(volName, true)
75+
}()
76+
err = mgr.removeVolume(volName, true) // why didn't this panic?
77+
if err == nil {
78+
t.Errorf("Expected error when removing nonexistent volume, but no error")
79+
}
80+
}
81+
82+
// T6: attach a volume that does not exist (should fail or panic)
83+
func TestAttachNonexistentVolume(t *testing.T) {
84+
mgr := setupMgr(t)
85+
86+
id, err := mgr.runContainerCuda("nonexistent_volume_t6")
87+
if id != "" && err != nil {
88+
t.Errorf("Expected error when removing nonexistent volume, but no error")
89+
}
90+
}
91+
92+
// T7: two containers attach to the same volume (should succeed in Docker, but test for your policy)
93+
func TestTwoContainersSameVolume(t *testing.T) {
94+
mgr := setupMgr(t)
95+
volName := "test_volume_t7"
96+
mgr.createVolume(volName)
97+
id1, err := mgr.runContainerCuda(volName)
98+
if err != nil {
99+
fmt.Errorf("Failed to start container: %v", err.Error())
100+
}
101+
id2, err := mgr.runContainerCuda(volName)
102+
if err != nil {
103+
fmt.Errorf("Failed to start container: %v", err.Error())
104+
}
105+
mgr.stopContainer(id1)
106+
mgr.removeContainer(id1)
107+
mgr.stopContainer(id2)
108+
mgr.removeContainer(id2)
109+
mgr.removeVolume(volName, true)
110+
}
111+
112+
// T8: two containers try to attach to the same volume at the same time (should succeed in Docker)
113+
func TestTwoContainersSameVolumeConcurrent(t *testing.T) {
114+
mgr := setupMgr(t)
115+
volName := "test_volume_t8"
116+
mgr.createVolume(volName)
117+
id1, err := mgr.runContainerCuda(volName)
118+
if err != nil {
119+
fmt.Errorf("Failed to start container: %v", err.Error())
120+
}
121+
id2, err2 := mgr.runContainerCuda(volName)
122+
if err2 != nil {
123+
fmt.Errorf("Failed to start container: %v", err2.Error())
124+
}
125+
mgr.stopContainer(id1)
126+
mgr.removeContainer(id1)
127+
mgr.stopContainer(id2)
128+
mgr.removeContainer(id2)
129+
mgr.removeVolume(volName, true)
130+
}
131+
132+
// T9: set a limit of 100 volumes (should fail on 101st if you enforce a limit)
133+
func TestVolumeLimit(t *testing.T) {
134+
mgr := setupMgr(t)
135+
limit := 100
136+
created := []string{}
137+
for i := 0; i < limit; i++ {
138+
name := "test_volume_t9_" + fmt.Sprint(i)
139+
mgr.createVolume(name)
140+
created = append(created, name)
141+
}
142+
name := "test_volume_fail"
143+
_, err := mgr.createVolume(name)
144+
if err == nil {
145+
fmt.Errorf("Volume limit not enforced")
146+
}
147+
148+
defer func() {
149+
for _, name := range created {
150+
mgr.removeVolume(name, true)
151+
}
152+
}()
153+
// why didn't you clean up the volumes
154+
// Try to create one more if you enforce a limit
155+
// If not enforced, this will succeed
156+
}
157+
158+
// T10: set a limit of 10 containers (should fail on 11th if you enforce a limit)
159+
func TestContainerLimit(t *testing.T) {
160+
mgr := setupMgr(t)
161+
volName := "test_volume_t10"
162+
mgr.createVolume(volName)
163+
ids := []string{}
164+
limit := 10
165+
for i := 0; i < limit; i++ {
166+
id, err := mgr.runContainerCuda(volName)
167+
if err != nil {
168+
fmt.Errorf("Failed to start container: %v", err.Error())
169+
}
170+
ids = append(ids, id)
171+
}
172+
_, err := mgr.runContainerCuda(volName)
173+
if err == nil {
174+
fmt.Errorf("Container limit not enforced")
175+
}
176+
defer func() {
177+
for _, id := range ids {
178+
mgr.stopContainer(id)
179+
mgr.removeContainer(id)
180+
}
181+
mgr.removeVolume(volName, true) // why didnt you clean up the containers?
182+
}()
183+
// Try to create one more if you enforce a limit
184+
// If not enforced, this will succeed
185+
}

0 commit comments

Comments
 (0)