Added validation for driver and executor shape details. (#211)

Author: lu-ohai

ads/jobs/builders/infrastructure/dataflow.py

@@ -42,6 +42,13 @@
 DEFAULT_SPARK_VERSION = "3.2.1"
 DEFAULT_NUM_EXECUTORS = 1
 DEFAULT_SHAPE = "VM.Standard.E3.Flex"
+DATAFLOW_SHAPE_FAMILY = [
+    "Standard.E3",
+    "Standard.E4",
+    "Standard3",
+    "Standard.A1",
+    "Standard2"
+]
 
 
 def conda_pack_name_to_dataflow_config(conda_uri):
@@ -889,11 +896,39 @@ def create(self, runtime: DataFlowRuntime, **kwargs) -> "DataFlow":
             raise ValueError(
                 "Compartment id is required. Specify compartment id via 'with_compartment_id()'."
             )
+        self._validate_shapes(payload)
         payload.pop("id", None)
         logger.debug(f"Creating a DataFlow Application with payload {payload}")
         self.df_app = DataFlowApp(**payload).create()
         self.with_id(self.df_app.id)
         return self
+    
+    @staticmethod
+    def _validate_shapes(payload: Dict):
+        if "executor_shape" not in payload:
+            payload["executor_shape"] = DEFAULT_SHAPE
+        if "driver_shape" not in payload:
+            payload["driver_shape"] = DEFAULT_SHAPE
+        executor_shape = payload["executor_shape"]
+        executor_shape_config = payload.get("executor_shape_config", {})
+        driver_shape = payload["driver_shape"]
+        driver_shape_config = payload.get("driver_shape_config", {})
+        same_shape_family = False
+        for shape in DATAFLOW_SHAPE_FAMILY:
+            if shape in executor_shape and shape in driver_shape:
+                same_shape_family = True
+                break
+        if not same_shape_family:
+            raise ValueError(
+                "`executor_shape` and `driver_shape` must be from the same shape family."
+            )
+        if (
+            (not executor_shape.endswith("Flex") and executor_shape_config) 
+            or (not driver_shape.endswith("Flex") and driver_shape_config)
+        ):
+            raise ValueError(
+                "Shape config is not required for non flex shape from user end."
+            )
 
     @staticmethod
     def _upload_file(local_path, bucket, overwrite=False):

tests/unitary/default_setup/jobs/test_jobs_dataflow.py

@@ -37,9 +37,9 @@
     arguments=["test-df"],
     compartment_id="ocid1.compartment.oc1..<unique_ocid>",
     display_name="test-df",
-    driver_shape="VM.Standard2.1",
+    driver_shape="VM.Standard.E4.Flex",
     driver_shape_config={"memory_in_gbs": 1, "ocpus": 16},
-    executor_shape="VM.Standard2.1",
+    executor_shape="VM.Standard.E4.Flex",
     executor_shape_config={"memory_in_gbs": 1, "ocpus": 16},
     file_uri="oci://test_bucket@test_namespace/test-dataflow/test-dataflow.py",
     num_executors=1,
@@ -124,7 +124,7 @@ def test_create_delete(self, mock_to_dict, mock_client):
                     df.lifecycle_state
                     == oci.data_flow.models.Application.LIFECYCLE_STATE_DELETED
                 )
-                assert len(df.to_yaml()) == 557
+                assert len(df.to_yaml()) == 567
 
     def test_create_df_app_with_default_display_name(
         self,
@@ -433,8 +433,8 @@ def test_create_from_id(self, mock_from_ocid):
         mock_from_ocid.return_value = Application(**SAMPLE_PAYLOAD)
         df = DataFlow.from_id("ocid1.datasciencejob.oc1.iad.<unique_ocid>")
         assert df.name == "test-df"
-        assert df.driver_shape == "VM.Standard2.1"
-        assert df.executor_shape == "VM.Standard2.1"
+        assert df.driver_shape == "VM.Standard.E4.Flex"
+        assert df.executor_shape == "VM.Standard.E4.Flex"
         assert df.private_endpoint_id == "test_private_endpoint"
 
         assert (
@@ -454,7 +454,7 @@ def test_create_from_id(self, mock_from_ocid):
     def test_to_and_from_dict(self, df):
         df_dict = df.to_dict()
         assert df_dict["spec"]["numExecutors"] == 2
-        assert df_dict["spec"]["driverShape"] == "VM.Standard2.1"
+        assert df_dict["spec"]["driverShape"] == "VM.Standard.E4.Flex"
         assert df_dict["spec"]["logsBucketUri"] == "oci://test_bucket@test_namespace/"
         assert df_dict["spec"]["privateEndpointId"] == "test_private_endpoint"
         assert df_dict["spec"]["driverShapeConfig"] == {"memoryInGBs": 1, "ocpus": 16}
@@ -482,6 +482,51 @@ def test_to_and_from_dict(self, df):
         assert df3_dict["spec"]["sparkVersion"] == "3.2.1"
         assert df3_dict["spec"]["numExecutors"] == 2
 
+    def test_shape_and_details(self, mock_to_dict, mock_client, df):
+        df.with_driver_shape(
+            "VM.Standard2.1"
+        ).with_executor_shape(
+            "VM.Standard.E4.Flex"
+        )
+
+        rt = (
+            DataFlowRuntime()
+                .with_script_uri(SAMPLE_PAYLOAD["file_uri"])
+                .with_archive_uri(SAMPLE_PAYLOAD["archive_uri"])
+                .with_custom_conda(
+                    "oci://my_bucket@my_namespace/conda_environments/cpu/PySpark 3.0 and Data Flow/5.0/pyspark30_p37_cpu_v5"
+                )
+                .with_overwrite(True)
+        )
+
+        with pytest.raises(
+            ValueError,
+            match="`executor_shape` and `driver_shape` must be from the same shape family."
+        ):
+            with patch.object(DataFlowApp, "client", mock_client):
+                with patch.object(DataFlowApp, "to_dict", mock_to_dict):
+                    df.create(rt)
+
+        df.with_driver_shape(
+            "VM.Standard2.1"
+        ).with_driver_shape_config(
+            memory_in_gbs=SAMPLE_PAYLOAD["driver_shape_config"]["memory_in_gbs"],
+            ocpus=SAMPLE_PAYLOAD["driver_shape_config"]["ocpus"],
+        ).with_executor_shape(
+            "VM.Standard2.16"
+        ).with_executor_shape_config(
+            memory_in_gbs=SAMPLE_PAYLOAD["executor_shape_config"]["memory_in_gbs"],
+            ocpus=SAMPLE_PAYLOAD["executor_shape_config"]["ocpus"],
+        )
+
+        with pytest.raises(
+            ValueError,
+            match="Shape config is not required for non flex shape from user end."
+        ):
+            with patch.object(DataFlowApp, "client", mock_client):
+                with patch.object(DataFlowApp, "to_dict", mock_to_dict):
+                    df.create(rt)
+            
 
 class TestDataFlowNotebookRuntime:
     @pytest.mark.skipif(