Add dockerRunOptions setting (#64)

* Add dockerRunOptions setting Signed-off-by: Jacob Tomlinson <[email protected]> * Add missing args Signed-off-by: Jacob Tomlinson <[email protected]> --------- Signed-off-by: Jacob Tomlinson <[email protected]>
NVIDIA · Jul 23, 2024 · 8f14a3e · 8f14a3e
1 parent c473896
commit 8f14a3e
Show file tree

Hide file tree

Showing 7 changed files with 33 additions and 6 deletions.
diff --git a/examples/gpu.yaml b/examples/gpu.yaml
@@ -0,0 +1,18 @@
+apiVersion: container-canary.nvidia.com/v1
+kind: Validator
+name: gpu
+description: A GPU example to show you can pass extra flags to Docker
+command:
+- "sleep"
+- "30"
+dockerRunOptions:
+- "--gpus"
+- "all"
+checks:
+  - name: nvidia-smi
+    description: 📦 Can run nvidia-smi
+    probe:
+      exec:
+        command:
+          - /usr/bin/nvidia-smi
+      initialDelaySeconds: 1
diff --git a/internal/apis/v1/types.go b/internal/apis/v1/types.go
@@ -48,6 +48,10 @@ type Validator struct {
 	// A command to run in the container
 	// +optional
 	Command []string
+
+	// Additional flags to pass to the docker CLI.
+	// +optional
+	DockerRunOptions []string  `yaml:"dockerRunOptions"`
 }
 
 type Check struct {

diff --git a/internal/config/load.go b/internal/config/load.go
@@ -26,7 +26,7 @@ import (
 	"path/filepath"
 
 	canaryv1 "github.com/nvidia/container-canary/internal/apis/v1"
-	"gopkg.in/yaml.v2"
+	yaml "gopkg.in/yaml.v2"
 )
 
 func LoadValidatorFromURL(url string) (*canaryv1.Validator, error) {

diff --git a/internal/container/container.go b/internal/container/container.go
@@ -44,7 +44,7 @@ type ContainerInterface interface {
 	Logs() (string, error)
 }
 
-func New(image string, env []v1.EnvVar, ports []v1.ServicePort, volumes []canaryv1.Volume, command []string) ContainerInterface {
+func New(image string, env []v1.EnvVar, ports []v1.ServicePort, volumes []canaryv1.Volume, command []string, dockerRunOptions []string) ContainerInterface {
 	name := fmt.Sprintf("%s%s", "canary-runner-", uuid.New().String()[:8])
-	return &DockerContainer{Name: name, Image: image, Command: command, Env: env, Ports: ports, Volumes: volumes}
+	return &DockerContainer{Name: name, Image: image, Command: command, Env: env, Ports: ports, Volumes: volumes, RunOptions: dockerRunOptions}
 }
diff --git a/internal/container/docker.go b/internal/container/docker.go
@@ -20,6 +20,7 @@ type DockerContainer struct {
 	Env        []v1.EnvVar
 	Ports      []v1.ServicePort
 	Volumes    []canaryv1.Volume
+	RunOptions []string
 	runCommand string
 }
 
@@ -46,6 +47,10 @@ func (c *DockerContainer) Start() error {
 		}
 	}
 
+	if len(c.RunOptions) > 0 {
+		commandArgs = append(commandArgs, c.RunOptions...)
+	}
+
 	commandArgs = append(commandArgs, c.Image)
 
 	if len(c.Command) > 0 {

diff --git a/internal/container/docker_test.go b/internal/container/docker_test.go
@@ -37,7 +37,7 @@ func TestDockerContainer(t *testing.T) {
 	volumes := []canaryv1.Volume{
 		{MountPath: "/foo"},
 	}
-	c := New("nginx", env, ports, volumes, nil)
+	c := New("nginx", env, ports, volumes, nil, nil)
 
 	err := c.Start()
 
@@ -69,7 +69,7 @@ func TestDockerContainer(t *testing.T) {
 	}
 }
 func TestDockerContainerRemoves(t *testing.T) {
-	c := New("nginx", nil, nil, nil, nil)
+	c := New("nginx", nil, nil, nil, nil, nil)
 
 	err := c.Start()
 	if err != nil {

diff --git a/internal/validator/validator.go b/internal/validator/validator.go
@@ -125,7 +125,7 @@ func loadConfig(filePath string) tea.Cmd {
 
 func startContainer(image string, validator *canaryv1.Validator) tea.Cmd {
 	return func() tea.Msg {
-		container := container.New(image, validator.Env, validator.Ports, validator.Volumes, validator.Command)
+		container := container.New(image, validator.Env, validator.Ports, validator.Volumes, validator.Command, validator.DockerRunOptions)
 		err := container.Start()
 		if err != nil {
 			return containerFailed{Error: err}