Updated pr.

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):
@@ -860,6 +867,15 @@ 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:
@@ -868,15 +884,22 @@ def create(self, runtime: DataFlowRuntime, **kwargs) -> "DataFlow":
         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()
-        self.with_id(self.df_app.id)
-        return self
+        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

@@ -448,7 +448,7 @@ 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"
+            "VM.Standard.E4.Flex"
         )
 
         rt = (
@@ -475,7 +475,7 @@ def test_shape_and_details(self, mock_to_dict, mock_client, df):
             memory_in_gbs=SAMPLE_PAYLOAD["driver_shape_config"]["memory_in_gbs"],
             ocpus=SAMPLE_PAYLOAD["driver_shape_config"]["ocpus"],
         ).with_executor_shape(
-            "VM.Standard2.1"
+            "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"],