From 100db97ffab11131a8ebe6fd45353e8f8248cf7e Mon Sep 17 00:00:00 2001
From: Martin <Junkern@users.noreply.github.com>
Date: Fri, 15 Jul 2022 11:15:28 +0200
Subject: [PATCH] feat: Add gpu flag to docker_container resource (#405)

* feat: Add gpu flag do docker_container resource.

* docs: Enhance gpus flag documentation.
---
 .../provider/resource_docker_container.go     |  6 ++++++
 .../resource_docker_container_funcs.go        | 20 +++++++++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/internal/provider/resource_docker_container.go b/internal/provider/resource_docker_container.go
index effdc507..dbb3d4be 100644
--- a/internal/provider/resource_docker_container.go
+++ b/internal/provider/resource_docker_container.go
@@ -973,6 +973,12 @@ func resourceDockerContainer() *schema.Resource {
 				Optional:    true,
 				ForceNew:    true,
 			},
+			"gpus": {
+				Type:        schema.TypeString,
+				Description: "GPU devices to add to the container. Currently, only the value `all` is supported. Passing any other value will result in unexpected behavior.",
+				Optional:    true,
+				ForceNew:    true,
+			},
 		},
 	}
 }
diff --git a/internal/provider/resource_docker_container_funcs.go b/internal/provider/resource_docker_container_funcs.go
index 988becbd..2d49474b 100644
--- a/internal/provider/resource_docker_container_funcs.go
+++ b/internal/provider/resource_docker_container_funcs.go
@@ -15,6 +15,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/docker/cli/opts"
 	"github.com/docker/docker/api/types"
 	"github.com/docker/docker/api/types/container"
 	"github.com/docker/docker/api/types/mount"
@@ -347,6 +348,18 @@ func resourceDockerContainerCreate(ctx context.Context, d *schema.ResourceData,
 	if v, ok := d.GetOk("group_add"); ok {
 		hostConfig.GroupAdd = stringSetToStringSlice(v.(*schema.Set))
 	}
+	if v, ok := d.GetOk("gpus"); ok {
+		if client.ClientVersion() >= "1.40" {
+			var gpu opts.GpuOpts
+			err := gpu.Set(v.(string))
+			if err != nil {
+				return diag.Errorf("Error setting gpus: %s", err)
+			}
+			hostConfig.DeviceRequests = gpu.Value()
+		} else {
+			log.Printf("[WARN] GPU support requires docker version 1.40 or higher")
+		}
+	}
 
 	init := d.Get("init").(bool)
 	hostConfig.Init = &init
@@ -730,6 +743,13 @@ func resourceDockerContainerRead(ctx context.Context, d *schema.ResourceData, me
 	d.Set("stop_signal", container.Config.StopSignal)
 	d.Set("stop_timeout", container.Config.StopTimeout)
 
+	if len(container.HostConfig.DeviceRequests) > 0 {
+		// TODO pass the original gpus property string back to the resource
+		// var gpuOpts opts.GpuOpts
+		// gpuOpts = opts.GpuOpts{container.HostConfig.DeviceRequests}
+		d.Set("gpus", "all")
+	}
+
 	return nil
 }