Added validation for driver and executor shape details.

Author: lu-ohai

ads/jobs/builders/infrastructure/dataflow.py

@@ -860,6 +860,18 @@ def create(self, runtime: DataFlowRuntime, **kwargs) -> "DataFlow":
             raise ValueError(
                 "Compartment id is required. Specify compartment id via 'with_compartment_id()'."
             )
+        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", {})
+        if executor_shape != driver_shape:
+            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.")
         payload.pop("id", None)
         logger.debug(f"Creating a DataFlow Application with payload {payload}")
         self.df_app = DataFlowApp(**payload).create()

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,
@@ -403,8 +403,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 (
@@ -424,7 +424,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}
@@ -444,6 +444,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.Standard2.8"
+        )
+
+        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.1"
+        ).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(