diff --git a/src/images/serve_image.go b/src/images/serve_image.go new file mode 100644 index 0000000..6084d2c --- /dev/null +++ b/src/images/serve_image.go @@ -0,0 +1,165 @@ +package main + +import ( + "context" + "fmt" + "io" + "os" + + "github.com/docker/docker/api/types/container" + "github.com/docker/docker/api/types/mount" + "github.com/docker/docker/api/types/volume" + "github.com/docker/docker/client" +) + +type ContainerMgr struct { + ctx context.Context + cli *client.Client + containerLimit int + volumeLimit int + containers map[string]struct{} + volumes map[string]struct{} +} + +func NewContainerMgr(client *client.Client, containerLimit, volumeLimit int) *ContainerMgr { + return &ContainerMgr{ + ctx: context.Background(), + cli: client, + containerLimit: containerLimit, + volumeLimit: volumeLimit, + containers: make(map[string]struct{}), + volumes: make(map[string]struct{}), + } +} + +func (mgr *ContainerMgr) stopContainer(containerID string) { + ctx := mgr.ctx + cli := mgr.cli + + err := cli.ContainerStop(ctx, containerID, container.StopOptions{}) + if err != nil { + panic(err) + } + +} + +func (mgr *ContainerMgr) removeContainer(containerID string) error { + ctx := mgr.ctx + cli := mgr.cli + err := cli.ContainerRemove(ctx, containerID, container.RemoveOptions{RemoveVolumes: true}) + if err != nil { + return err + } + delete(mgr.containers, containerID) + return nil +} + +func (mgr *ContainerMgr) createVolume(volumeName string) (volume.Volume, error) { + if len(mgr.volumes) >= mgr.volumeLimit { + return volume.Volume{}, fmt.Errorf("volume limit reached") + } + ctx := mgr.ctx + cli := mgr.cli + + vol, err := cli.VolumeCreate(ctx, volume.CreateOptions{ + Name: volumeName, // You can leave this empty for a random name + }) + if err != nil { + return volume.Volume{}, err + } + mgr.volumes[vol.Name] = struct{}{} + return vol, nil +} + +func (mgr *ContainerMgr) removeVolume(volumeName string, force bool) error { + ctx := mgr.ctx + cli := mgr.cli + + vols, _ := cli.VolumeList(ctx, volume.ListOptions{}) + found := false + for _, v := range vols.Volumes { + if v.Name == volumeName { + found = true + break + } + } + if !found { + return fmt.Errorf("volume %s does not exist", volumeName) + } + + err := cli.VolumeRemove(ctx, volumeName, force) + if err != nil { + return err + } + delete(mgr.volumes, volumeName) + return nil +} + +func (mgr *ContainerMgr) runContainerCuda(volumeName string) (string, error) { + if len(mgr.containers) >= mgr.containerLimit { + return "", fmt.Errorf("container limit reached") + } + ctx := mgr.ctx + cli := mgr.cli + + resp, err := cli.ContainerCreate(ctx, &container.Config{ + Image: "pytorch-cuda", + Cmd: []string{"sleep", "1000"}, + }, &container.HostConfig{ + Runtime: "nvidia", + Mounts: []mount.Mount{ + { + Type: mount.TypeVolume, + Source: volumeName, + Target: "/data", + }, + }, + }, nil, nil, "") + vols, _ := cli.VolumeList(ctx, volume.ListOptions{}) + found := false + for _, v := range vols.Volumes { + if v.Name == volumeName { + found = true + break + } + } + if !found { + return "", fmt.Errorf("volume %s does not exist", volumeName) + } + if err != nil { + return "", err + } + mgr.containers[resp.ID] = struct{}{} + if err := cli.ContainerStart(ctx, resp.ID, container.StartOptions{}); err != nil { + return "", err + } + + out, err := cli.ContainerLogs(ctx, resp.ID, container.LogsOptions{ShowStdout: true}) + if err != nil { + panic(err) + } + + io.Copy(os.Stdout, out) + return resp.ID, nil +} + +func main() { + cli, err := client.NewClientWithOpts(client.FromEnv) + if err != nil { + panic(err) + } + + // Create a Docker volume + containerMgr := NewContainerMgr(cli, 10, 10) + volumeName := "my_volume1" + + containerMgr.createVolume(volumeName) + id, err := containerMgr.runContainerCuda(volumeName) + if err != nil { + fmt.Errorf("Failed to start container: %v", err.Error()) + } + containerMgr.stopContainer(id) + containerMgr.removeContainer(id) + containerMgr.removeVolume(volumeName, true) + +} diff --git a/src/images/serve_image_test.go b/src/images/serve_image_test.go new file mode 100644 index 0000000..d4a77db --- /dev/null +++ b/src/images/serve_image_test.go @@ -0,0 +1,219 @@ +package main + +import ( + "fmt" + "testing" + + "github.com/docker/docker/api/types/volume" + "github.com/docker/docker/client" +) + +func setupMgr(t *testing.T) *ContainerMgr { + cli, err := client.NewClientWithOpts(client.FromEnv) + if err != nil { + t.Fatalf("Failed to create Docker client: %v", err) + } + return NewContainerMgr(cli, 10, 100) +} + +// T1: create a volume, check exists, delete, check not exists +func TestCreateDeleteVolume(t *testing.T) { + mgr := setupMgr(t) + volName := "test_volume_t1" + mgr.createVolume(volName) + // vols, _ := mgr.cli.VolumeList(mgr.ctx, *opts*/ {}) + vols, _ := mgr.cli.VolumeList(mgr.ctx, volume.ListOptions{}) + found := false + for _, v := range vols.Volumes { + if v.Name == volName { + found = true + break + } + } + if !found { + t.Errorf("Volume %s not found after creation", volName) + } + mgr.removeVolume(volName, true) + // vols, _ = mgr.cli.VolumeList(mgr.ctx, /*opts*/ {}) + vols, _ = mgr.cli.VolumeList(mgr.ctx, volume.ListOptions{}) + for _, v := range vols.Volumes { + if v.Name == volName { + t.Errorf("Volume %s still exists after deletion", volName) + } + } +} + +// T2: create volume, start container, attach, write, stop, start, check persistence, cleanup +func TestVolumePersistence(t *testing.T) { + mgr := setupMgr(t) + volName := "test_volume_t2" + mgr.createVolume(volName) + containerID, err := mgr.runContainerCuda(volName) + if err != nil { + fmt.Errorf("Failed to start container: %v", err.Error()) + } + // Write to volume (you'd need to exec into container or mount and write a file) + // For example, use mgr.execInContainer(containerID, "sh", "-c", "echo hello > /data/test.txt") + // Stop and start container-p + mgr.stopContainer(containerID) + // mgr.startContainer(containerID) + // Check file exists (again, exec into container and check) + // Cleanup + mgr.stopContainer(containerID) + mgr.removeContainer(containerID) + mgr.removeVolume(volName, true) +} + +// T3: create a volume with same name twice (should not fail) +func TestCreateVolumeTwice(t *testing.T) { + mgr := setupMgr(t) + volName := "test_volume_t3" + mgr.createVolume(volName) + defer mgr.removeVolume(volName, true) + mgr.createVolume(volName) // Should not fail +} + +// T4: remove volume that doesn't exist (should fail or panic) +func TestRemoveNonexistentVolume(t *testing.T) { + mgr := setupMgr(t) + // defer func() { + // if r := recover(); r == nil { + // t.Errorf("Expected panic when removing nonexistent volume, but did not panic") + // } + // }() + err := mgr.removeVolume("nonexistent_volume_t4", true) // Maybe this function never panics + if err == nil { + t.Errorf("Expected error when removing nonexistent volume, but no error") + } +} + +// T5: remove volume in use (should fail or panic) +func TestRemoveVolumeInUse(t *testing.T) { + mgr := setupMgr(t) + volName := "test_volume_t5" + mgr.createVolume(volName) + containerID, err := mgr.runContainerCuda(volName) + if err != nil { + fmt.Errorf("Failed to start container: %v", err.Error()) + } + defer func() { + mgr.stopContainer(containerID) + mgr.removeContainer(containerID) + mgr.removeVolume(volName, true) + }() + err = mgr.removeVolume(volName, true) // why didn't this panic? + if err == nil { + t.Errorf("Expected error when removing nonexistent volume, but no error") + } +} + +// T6: attach a volume that does not exist (should fail or panic) +func TestAttachNonexistentVolume(t *testing.T) { + mgr := setupMgr(t) + // defer func() { + // if r := recover(); r == nil { + // t.Errorf("Expected panic when attaching nonexistent volume, but did not panic") + // } + // }() + + id, err := mgr.runContainerCuda("nonexistent_volume_t6") // why did this one panic/work + // fmt.Printf("ALLO: %v, %v", id, err) + if id != "" && err != nil { + t.Errorf("Expected error when removing nonexistent volume, but no error") + } +} + +// T7: two containers attach to the same volume (should succeed in Docker, but test for your policy) +func TestTwoContainersSameVolume(t *testing.T) { + mgr := setupMgr(t) + volName := "test_volume_t7" + mgr.createVolume(volName) + id1, err := mgr.runContainerCuda(volName) + if err != nil { + fmt.Errorf("Failed to start container: %v", err.Error()) + } + id2, err := mgr.runContainerCuda(volName) + if err != nil { + fmt.Errorf("Failed to start container: %v", err.Error()) + } + mgr.stopContainer(id1) + mgr.removeContainer(id1) + mgr.stopContainer(id2) + mgr.removeContainer(id2) + mgr.removeVolume(volName, true) +} + +// T8: two containers try to attach to the same volume at the same time (should succeed in Docker) +func TestTwoContainersSameVolumeConcurrent(t *testing.T) { + mgr := setupMgr(t) + volName := "test_volume_t8" + mgr.createVolume(volName) + id1, err := mgr.runContainerCuda(volName) + if err != nil { + fmt.Errorf("Failed to start container: %v", err.Error()) + } + id2, err2 := mgr.runContainerCuda(volName) + if err2 != nil { + fmt.Errorf("Failed to start container: %v", err2.Error()) + } + mgr.stopContainer(id1) + mgr.removeContainer(id1) + mgr.stopContainer(id2) + mgr.removeContainer(id2) + mgr.removeVolume(volName, true) +} + +// T9: set a limit of 100 volumes (should fail on 101st if you enforce a limit) +func TestVolumeLimit(t *testing.T) { + mgr := setupMgr(t) + limit := 100 + created := []string{} + for i := 0; i < limit; i++ { + name := "test_volume_t9_" + fmt.Sprint(i) + mgr.createVolume(name) + created = append(created, name) + } + name := "test_volume_fail" + _, err := mgr.createVolume(name) + if err == nil { + fmt.Errorf("Volume limit not enforced") + } + + defer func() { + for _, name := range created { + mgr.removeVolume(name, true) + } + }() + // why didn't you clean up the volumes + // Try to create one more if you enforce a limit + // If not enforced, this will succeed +} + +// T10: set a limit of 10 containers (should fail on 11th if you enforce a limit) +func TestContainerLimit(t *testing.T) { + mgr := setupMgr(t) + volName := "test_volume_t10" + mgr.createVolume(volName) + ids := []string{} + limit := 10 + for i := 0; i < limit; i++ { + id, err := mgr.runContainerCuda(volName) + if err != nil { + fmt.Errorf("Failed to start container: %v", err.Error()) + } + ids = append(ids, id) + } + _, err := mgr.runContainerCuda(volName) + if err == nil { + fmt.Errorf("Container limit not enforced") + } + defer func() { + for _, id := range ids { + mgr.stopContainer(id) + mgr.removeContainer(id) + } + mgr.removeVolume(volName, true) // why didnt you clean up the containers? + }() + // Try to create one more if you enforce a limit + // If not enforced, this will succeed +}