f/aws_sagemaker_notebook_instance: Volume size

aws/resource_aws_sagemaker_notebook_instance.go

@@ -46,6 +46,12 @@ func resourceAwsSagemakerNotebookInstance() *schema.Resource {
 				Required: true,
 			},
 
+			"volume_size": {
+				Type:     schema.TypeInt,
+				Optional: true,
+				Default:  5,
+			},
+
 			"subnet_id": {
 				Type:     schema.TypeString,
 				Optional: true,
@@ -123,6 +129,10 @@ func resourceAwsSagemakerNotebookInstanceCreate(d *schema.ResourceData, meta int
 		createOpts.SubnetId = aws.String(s.(string))
 	}
 
+	if v, ok := d.GetOk("volume_size"); ok {
+		createOpts.VolumeSizeInGB = aws.Int64(int64(v.(int)))
+	}
+
 	if k, ok := d.GetOk("kms_key_id"); ok {
 		createOpts.KmsKeyId = aws.String(k.(string))
 	}
@@ -200,6 +210,10 @@ func resourceAwsSagemakerNotebookInstanceRead(d *schema.ResourceData, meta inter
 		return fmt.Errorf("error setting kms_key_id for sagemaker notebook instance (%s): %s", d.Id(), err)
 	}
 
+	if err := d.Set("volume_size", notebookInstance.VolumeSizeInGB); err != nil {
+		return fmt.Errorf("error setting volume_size for sagemaker notebook instance (%s): %s", d.Id(), err)
+	}
+
 	if err := d.Set("lifecycle_config_name", notebookInstance.NotebookInstanceLifecycleConfigName); err != nil {
 		return fmt.Errorf("error setting lifecycle_config_name for sagemaker notebook instance (%s): %s", d.Id(), err)
 	}
@@ -256,6 +270,11 @@ func resourceAwsSagemakerNotebookInstanceUpdate(d *schema.ResourceData, meta int
 		hasChanged = true
 	}
 
+	if d.HasChange("volume_size") {
+		updateOpts.VolumeSizeInGB = aws.Int64(int64(d.Get("volume_size").(int)))
+		hasChanged = true
+	}
+
 	if hasChanged {
 
 		// Stop notebook

aws/resource_aws_sagemaker_notebook_instance_test.go

@@ -96,6 +96,8 @@ func TestAccAWSSagemakerNotebookInstance_basic(t *testing.T) {
 
 					resource.TestCheckResourceAttr(
 						"aws_sagemaker_notebook_instance.foo", "name", notebookName),
+					resource.TestCheckResourceAttr(
+						"aws_sagemaker_notebook_instance.foo", "volume_size", "5"),
 				),
 			},
 			{
@@ -144,6 +146,43 @@ func TestAccAWSSagemakerNotebookInstance_update(t *testing.T) {
 	})
 }
 
+func TestAccAWSSagemakerNotebookInstance_volumesize(t *testing.T) {
+	var notebook sagemaker.DescribeNotebookInstanceOutput
+	notebookName := resource.PrefixedUniqueId(sagemakerTestAccSagemakerNotebookInstanceResourceNamePrefix)
+	var resourceName = "aws_sagemaker_notebook_instance.foo"
+	resource.ParallelTest(t, resource.TestCase{
+		PreCheck:     func() { testAccPreCheck(t) },
+		Providers:    testAccProviders,
+		CheckDestroy: testAccCheckAWSSagemakerNotebookInstanceDestroy,
+		Steps: []resource.TestStep{
+			{
+				Config: testAccAWSSagemakerNotebookInstanceConfigVolume(notebookName),
+				Check: resource.ComposeTestCheckFunc(
+					testAccCheckAWSSagemakerNotebookInstanceExists(resourceName, &notebook),
+
+					resource.TestCheckResourceAttr(
+						"aws_sagemaker_notebook_instance.foo", "volume_size", "5"),
+				),
+			},
+
+			{
+				Config: testAccAWSSagemakerNotebookInstanceUpdateConfig(notebookName),
+				Check: resource.ComposeTestCheckFunc(
+					testAccCheckAWSSagemakerNotebookInstanceExists("aws_sagemaker_notebook_instance.foo", &notebook),
+
+					resource.TestCheckResourceAttr(
+						"aws_sagemaker_notebook_instance.foo", "volume_size", "8"),
+				),
+			},
+			{
+				ResourceName:      resourceName,
+				ImportState:       true,
+				ImportStateVerify: true,
+			},
+		},
+	})
+}
+
 func TestAccAWSSagemakerNotebookInstance_LifecycleConfigName(t *testing.T) {
 	var notebook sagemaker.DescribeNotebookInstanceOutput
 	rName := resource.PrefixedUniqueId(sagemakerTestAccSagemakerNotebookInstanceResourceNamePrefix)
@@ -466,12 +505,41 @@ data "aws_iam_policy_document" "assume_role" {
 `, notebookName, notebookName)
 }
 
+func testAccAWSSagemakerNotebookInstanceConfigVolume(notebookName string) string {
+	return fmt.Sprintf(`
+resource "aws_sagemaker_notebook_instance" "foo" {
+  name          = "%s"
+  role_arn      = aws_iam_role.foo.arn
+  instance_type = "ml.t2.medium"
+  volume_size   = "5"
+}
+
+resource "aws_iam_role" "foo" {
+  name               = "%s"
+  path               = "/"
+  assume_role_policy = data.aws_iam_policy_document.assume_role.json
+}
+
+data "aws_iam_policy_document" "assume_role" {
+  statement {
+    actions = ["sts:AssumeRole"]
+
+    principals {
+      type        = "Service"
+      identifiers = ["sagemaker.amazonaws.com"]
+    }
+  }
+}
+`, notebookName, notebookName)
+}
+
 func testAccAWSSagemakerNotebookInstanceUpdateConfig(notebookName string) string {
 	return fmt.Sprintf(`
 resource "aws_sagemaker_notebook_instance" "foo" {
   name          = "%s"
   role_arn      = aws_iam_role.foo.arn
   instance_type = "ml.m4.xlarge"
+  volume_size   = "8"
 }
 
 resource "aws_iam_role" "foo" {

website/docs/r/sagemaker_notebook_instance.html.markdown

@@ -33,6 +33,7 @@ The following arguments are supported:
 * `name` - (Required) The name of the notebook instance (must be unique).
 * `role_arn` - (Required) The ARN of the IAM role to be used by the notebook instance which allows SageMaker to call other services on your behalf.
 * `instance_type` - (Required) The name of ML compute instance type.
+* `volume_size` - (Optional) The size, in GB, of the ML storage volume to attach to the notebook instance. The default value is 5 GB.
 * `subnet_id` - (Optional) The VPC subnet ID.
 * `security_groups` - (Optional) The associated security groups.
 * `kms_key_id` - (Optional) The AWS Key Management Service (AWS KMS) key that Amazon SageMaker uses to encrypt the model artifacts at rest using Amazon S3 server-side encryption.