fix: check optional keyword before accessing (#1911)

* fix: check optional keyword before accessing

* refactor attach method

Co-authored-by: Chuyang Deng <[email protected]>
Co-authored-by: icywang86rui <[email protected]>
Co-authored-by: Ajay Karpur <[email protected]>

Author: 4 people

src/sagemaker/model_monitor/model_monitoring.py

@@ -641,46 +641,32 @@ def attach(cls, monitor_schedule_name, sagemaker_session=None):
             monitoring_schedule_name=monitor_schedule_name
         )
 
-        role = schedule_desc["MonitoringScheduleConfig"]["MonitoringJobDefinition"]["RoleArn"]
-        image_uri = schedule_desc["MonitoringScheduleConfig"]["MonitoringJobDefinition"][
-            "MonitoringAppSpecification"
-        ]["ImageUri"]
-        instance_count = schedule_desc["MonitoringScheduleConfig"]["MonitoringJobDefinition"][
-            "MonitoringResources"
-        ]["ClusterConfig"]["InstanceCount"]
-        instance_type = schedule_desc["MonitoringScheduleConfig"]["MonitoringJobDefinition"][
-            "MonitoringResources"
-        ]["ClusterConfig"]["InstanceType"]
-        entrypoint = schedule_desc["MonitoringScheduleConfig"]["MonitoringJobDefinition"][
-            "MonitoringAppSpecification"
-        ].get("ContainerEntrypoint")
-        volume_size_in_gb = schedule_desc["MonitoringScheduleConfig"]["MonitoringJobDefinition"][
-            "MonitoringResources"
-        ]["ClusterConfig"]["VolumeSizeInGB"]
-        volume_kms_key = schedule_desc["MonitoringScheduleConfig"]["MonitoringJobDefinition"][
-            "MonitoringResources"
-        ]["ClusterConfig"].get("VolumeKmsKeyId")
-        output_kms_key = schedule_desc["MonitoringScheduleConfig"]["MonitoringJobDefinition"][
-            "MonitoringOutputConfig"
-        ].get("KmsKeyId")
+        monitoring_job_definition = schedule_desc["MonitoringScheduleConfig"][
+            "MonitoringJobDefinition"
+        ]
+        role = monitoring_job_definition["RoleArn"]
+        image_uri = monitoring_job_definition["MonitoringAppSpecification"].get("ImageUri")
+        cluster_config = monitoring_job_definition["MonitoringResources"]["ClusterConfig"]
+        instance_count = cluster_config.get("InstanceCount")
+        instance_type = cluster_config["InstanceType"]
+        volume_size_in_gb = cluster_config["VolumeSizeInGB"]
+        volume_kms_key = cluster_config.get("VolumeKmsKeyId")
+        entrypoint = monitoring_job_definition["MonitoringAppSpecification"].get(
+            "ContainerEntrypoint"
+        )
+        output_kms_key = monitoring_job_definition["MonitoringOutputConfig"].get("KmsKeyId")
+        network_config_dict = monitoring_job_definition.get("NetworkConfig")
 
         max_runtime_in_seconds = None
-        if schedule_desc["MonitoringScheduleConfig"]["MonitoringJobDefinition"].get(
-            "StoppingCondition"
-        ):
-            max_runtime_in_seconds = schedule_desc["MonitoringScheduleConfig"][
-                "MonitoringJobDefinition"
-            ]["StoppingCondition"].get("MaxRuntimeInSeconds")
+        stopping_condition = monitoring_job_definition.get("StoppingCondition")
+        if stopping_condition:
+            max_runtime_in_seconds = stopping_condition.get("MaxRuntimeInSeconds")
 
-        env = schedule_desc["MonitoringScheduleConfig"]["MonitoringJobDefinition"]["Environment"]
-
-        network_config_dict = schedule_desc["MonitoringScheduleConfig"][
-            "MonitoringJobDefinition"
-        ].get("NetworkConfig")
+        env = monitoring_job_definition.get("Environment", None)
 
-        vpc_config = schedule_desc["MonitoringScheduleConfig"]["MonitoringJobDefinition"][
-            "NetworkConfig"
-        ].get("VpcConfig")
+        vpc_config = None
+        if network_config_dict:
+            vpc_config = network_config_dict.get("VpcConfig")
 
         security_group_ids = None
         if vpc_config is not None:
@@ -690,6 +676,7 @@ def attach(cls, monitor_schedule_name, sagemaker_session=None):
         if vpc_config is not None:
             subnets = vpc_config["Subnets"]
 
+        network_config = None
         if network_config_dict:
             network_config = NetworkConfig(
                 enable_network_isolation=network_config_dict["EnableNetworkIsolation"],

tests/unit/sagemaker/monitor/test_model_monitoring.py

@@ -63,6 +63,31 @@
 )
 "encrypt_inter_container_traffic=None when creating your NetworkConfig object."
 
+MONITORING_SCHEDULE_DESC = {
+    "MonitoringScheduleArn": "arn:aws:monitoring-schedule",
+    "MonitoringScheduleName": "my-monitoring-schedule",
+    "MonitoringScheduleConfig": {
+        "MonitoringJobDefinition": {
+            "MonitoringOutputConfig": {},
+            "MonitoringResources": {
+                "ClusterConfig": {
+                    "InstanceCount": 1,
+                    "InstanceType": "ml.t3.medium",
+                    "VolumeSizeInGB": 8,
+                }
+            },
+            "MonitoringAppSpecification": {
+                "ImageUri": "image-uri",
+                "ContainerEntrypoint": [
+                    "entrypoint.py",
+                ],
+            },
+            "RoleArn": ROLE,
+        }
+    },
+    "EndpointName": "my-endpoint",
+}
+
 
 # TODO-reinvent-2019: Continue to flesh these out.
 @pytest.fixture()
@@ -80,6 +105,9 @@ def sagemaker_session():
         name="upload_data", return_value="mocked_s3_uri_from_upload_data"
     )
     session_mock.download_data = Mock(name="download_data")
+    session_mock.describe_monitoring_schedule = Mock(
+        name="describe_monitoring_schedule", return_value=MONITORING_SCHEDULE_DESC
+    )
     return session_mock
 
 
@@ -153,6 +181,17 @@ def test_default_model_monitor_with_invalid_network_config(sagemaker_session):
     assert INTER_CONTAINER_ENCRYPTION_EXCEPTION_MSG in str(exception.value)
 
 
+def test_model_monitor_without_network_config(sagemaker_session):
+    my_model_monitor = ModelMonitor(
+        role=ROLE,
+        image_uri=CUSTOM_IMAGE_URI,
+        sagemaker_session=sagemaker_session,
+    )
+    model_monitor_schedule_name = "model-monitoring-without-network-config"
+    attached = my_model_monitor.attach(model_monitor_schedule_name, sagemaker_session)
+    assert attached.network_config is None
+
+
 def test_model_monitor_with_invalid_network_config(sagemaker_session):
     invalid_network_config = NetworkConfig(encrypt_inter_container_traffic=False)
     my_model_monitor = ModelMonitor(