Add cpuset_mems to docker driver.

Signed-off-by: Shishir Mahajan <smahajan@roblox.com>

Author: shishir-a412ed

drivers/docker/config.go

@@ -341,6 +341,7 @@ var (
 		"cap_drop":       hclspec.NewAttr("cap_drop", "list(string)", false),
 		"command":        hclspec.NewAttr("command", "string", false),
 		"cpuset_cpus":    hclspec.NewAttr("cpuset_cpus", "string", false),
+		"cpuset_mems":    hclspec.NewAttr("cpuset_mems", "string", false),
 		"cpu_hard_limit": hclspec.NewAttr("cpu_hard_limit", "bool", false),
 		"cpu_cfs_period": hclspec.NewDefault(
 			hclspec.NewAttr("cpu_cfs_period", "number", false),
@@ -433,6 +434,7 @@ type TaskConfig struct {
 	CPUCFSPeriod      int64              `codec:"cpu_cfs_period"`
 	CPUHardLimit      bool               `codec:"cpu_hard_limit"`
 	CPUSetCPUs        string             `codec:"cpuset_cpus"`
+	CPUSetMEMs        string             `codec:"cpuset_mems"`
 	Devices           []DockerDevice     `codec:"devices"`
 	DNSSearchDomains  []string           `codec:"dns_search_domains"`
 	DNSOptions        []string           `codec:"dns_options"`

drivers/docker/driver.go

@@ -888,6 +888,13 @@ func (d *Driver) createContainerConfig(task *drivers.TaskConfig, driverConfig *T
 		hostConfig.CPUSetCPUs = driverConfig.CPUSetCPUs
 	}
 
+	// This translates to docker create/run --cpuset-mems option.
+	// --cpuset-mems is the list of memory nodes on which processes
+	// in this cpuset are allowed to allocate memory.
+	if driverConfig.CPUSetMEMs != "" {
+		hostConfig.CPUSetMEMs = driverConfig.CPUSetMEMs
+	}
+
 	// Enable tini (docker-init) init system.
 	if driverConfig.Init {
 		hostConfig.Init = driverConfig.Init

drivers/docker/driver_test.go

@@ -1553,6 +1553,28 @@ func TestDockerDriver_CPUSetCPUs(t *testing.T) {
 	}
 }
 
+func TestDockerDriver_CPUSetMEMs(t *testing.T) {
+	ci.Parallel(t)
+	testutil.DockerCompatible(t)
+	if runtime.GOOS == "windows" {
+		t.Skip("Windows does not support cpuset_mems")
+	}
+
+	task, cfg, _ := dockerTask(t)
+
+	cfg.CPUSetMEMs = "0"
+	require.NoError(t, task.EncodeConcreteDriverConfig(cfg))
+
+	client, d, handle, cleanup := dockerSetup(t, task, nil)
+	defer cleanup()
+	require.NoError(t, d.WaitUntilStarted(task.ID, 5*time.Second))
+
+	container, err := client.InspectContainer(handle.containerID)
+	require.NoError(t, err)
+
+	require.Equal(t, cfg.CPUSetMEMs, container.HostConfig.CPUSetMEMs)
+}
+
 func TestDockerDriver_MemoryHardLimit(t *testing.T) {
 	ci.Parallel(t)
 	testutil.DockerCompatible(t)

website/content/docs/drivers/docker.mdx

@@ -84,7 +84,7 @@ The `docker` driver supports the following configuration in the job spec. Only
   }
   ```
 
-- `cpuset_cpus` <sup>Beta</sup> - (Optional) CPUs in which to allow execution
+- `cpuset_cpus` - (Optional) CPUs in which to allow execution
   (0-3, 0,1). Limit the specific CPUs or cores a container can use. A
   comma-separated list or hyphen-separated range of CPUs a container can use, if
   you have more than one CPU. The first CPU is numbered 0. A valid value might
@@ -100,6 +100,22 @@ config {
 }
 ```
 
+- `cpuset_mems` - (Optional) MEMs in which to allow execution
+  (0-3, 0,1). Limit the specific memory nodes a container process is allowed
+  to allocate memory. A comma-separated list or hyphen-separated range of memory
+  nodes a container can use, if you have more than one memory node. The first
+  memory node is numbered 0. A valid value might be 0-3 (to use the first, second,
+  third, and fourth memory node) or 1,3 (to use the second and fourth memory node).
+
+Note: `cpuset_mems` pins the workload to the memory nodes but doesn't give the workload
+exclusive access to those memory nodes.
+
+```hcl
+config {
+  cpuset_mems = "0-3"
+}
+```
+
 - `dns_search_domains` - (Optional) A list of DNS search domains for
   the container to use. If you are using bridge networking mode with a
   `network` block in the task group, you must set all DNS options in