Added flex shape for opctl config (#256)

Author: lu-ohai

ads/model/deployment/model_deployment.py

@@ -29,6 +29,11 @@
 from ads.model.common.utils import _is_json_serializable
 from ads.model.deployment.common.utils import send_request
 from ads.model.deployment.model_deployment_infrastructure import (
+    DEFAULT_BANDWIDTH_MBPS,
+    DEFAULT_REPLICA,
+    DEFAULT_SHAPE_NAME,
+    DEFAULT_OCPUS,
+    DEFAULT_MEMORY_IN_GBS,
     MODEL_DEPLOYMENT_INFRASTRUCTURE_TYPE,
     ModelDeploymentInfrastructure,
 )
@@ -64,12 +69,6 @@
 MODEL_DEPLOYMENT_TYPE = "modelDeployment"
 MODEL_DEPLOYMENT_INFERENCE_SERVER_TRITON = "TRITON"
 
-MODEL_DEPLOYMENT_INSTANCE_SHAPE = "VM.Standard.E4.Flex"
-MODEL_DEPLOYMENT_INSTANCE_OCPUS = 1
-MODEL_DEPLOYMENT_INSTANCE_MEMORY_IN_GBS = 16
-MODEL_DEPLOYMENT_INSTANCE_COUNT = 1
-MODEL_DEPLOYMENT_BANDWIDTH_MBPS = 10
-
 MODEL_DEPLOYMENT_RUNTIMES = {
     ModelDeploymentRuntimeType.CONDA: ModelDeploymentCondaRuntime,
     ModelDeploymentRuntimeType.CONTAINER: ModelDeploymentContainerRuntime,
@@ -1601,7 +1600,7 @@ def _build_model_deployment_configuration_details(self) -> Dict:
 
         instance_configuration = {
             infrastructure.CONST_INSTANCE_SHAPE_NAME: infrastructure.shape_name
-            or MODEL_DEPLOYMENT_INSTANCE_SHAPE,
+            or DEFAULT_SHAPE_NAME,
         }
 
         if instance_configuration[infrastructure.CONST_INSTANCE_SHAPE_NAME].endswith(
@@ -1613,14 +1612,14 @@ def _build_model_deployment_configuration_details(self) -> Dict:
                 infrastructure.CONST_OCPUS: infrastructure.shape_config_details.get(
                     "ocpus", None
                 )
-                or MODEL_DEPLOYMENT_INSTANCE_OCPUS,
+                or DEFAULT_OCPUS,
                 infrastructure.CONST_MEMORY_IN_GBS: infrastructure.shape_config_details.get(
                     "memory_in_gbs", None
                 )
                 or infrastructure.shape_config_details.get(
                     "memoryInGBs", None
                 )
-                or MODEL_DEPLOYMENT_INSTANCE_MEMORY_IN_GBS,
+                or DEFAULT_MEMORY_IN_GBS,
             }
 
         if infrastructure.subnet_id:
@@ -1629,7 +1628,7 @@ def _build_model_deployment_configuration_details(self) -> Dict:
         scaling_policy = {
             infrastructure.CONST_POLICY_TYPE: "FIXED_SIZE",
             infrastructure.CONST_INSTANCE_COUNT: infrastructure.replica
-            or MODEL_DEPLOYMENT_INSTANCE_COUNT,
+            or DEFAULT_REPLICA,
         }
 
         if not runtime.model_uri:
@@ -1660,7 +1659,7 @@ def _build_model_deployment_configuration_details(self) -> Dict:
 
         model_configuration_details = {
             infrastructure.CONST_BANDWIDTH_MBPS: infrastructure.bandwidth_mbps
-            or MODEL_DEPLOYMENT_BANDWIDTH_MBPS,
+            or DEFAULT_BANDWIDTH_MBPS,
             infrastructure.CONST_INSTANCE_CONFIG: instance_configuration,
             runtime.CONST_MODEL_ID: model_id,
             infrastructure.CONST_SCALING_POLICY: scaling_policy,

ads/model/deployment/model_deployment_infrastructure.py

@@ -22,7 +22,9 @@
 DEFAULT_BANDWIDTH_MBPS = 10
 DEFAULT_WEB_CONCURRENCY = 10
 DEFAULT_REPLICA = 1
-DEFAULT_SHAPE_NAME = "VM.Standard.E2.4"
+DEFAULT_SHAPE_NAME = "VM.Standard.E4.Flex"
+DEFAULT_OCPUS = 1
+DEFAULT_MEMORY_IN_GBS = 16
 
 logger = logging.getLogger(__name__)
 
@@ -625,4 +627,8 @@ def init(self) -> "ModelDeploymentInfrastructure":
             .with_web_concurrency(self.web_concurrency or DEFAULT_WEB_CONCURRENCY)
             .with_replica(self.replica or DEFAULT_REPLICA)
             .with_shape_name(self.shape_name or DEFAULT_SHAPE_NAME)
+            .with_shape_config_details(
+                ocpus=self.shape_config_details.get(self.CONST_OCPUS, DEFAULT_OCPUS),
+                memory_in_gbs=self.shape_config_details.get(self.CONST_MEMORY_IN_GBS, DEFAULT_MEMORY_IN_GBS)
+            )
         )

ads/opctl/cmds.py

@@ -593,6 +593,8 @@ def configure() -> None:
             ("log_id", ""),
             ("docker_registry", ""),
             ("conda_pack_os_prefix", "in the format oci://<bucket>@<namespace>/<path>"),
+            ("memory_in_gbs", ""),
+            ("ocpus", "")
         ]
         _set_service_configurations(
             ADS_JOBS_CONFIG_FILE_NAME,
@@ -619,6 +621,10 @@ def configure() -> None:
             ("num_executors", ""),
             ("spark_version", ""),
             ("archive_bucket", "in the format oci://<bucket>@<namespace>/<path>"),
+            ("driver_shape_memory_in_gbs", ""),
+            ("driver_shape_ocpus", ""),
+            ("executor_shape_memory_in_gbs", ""),
+            ("executor_shape_ocpus", "")
         ]
         _set_service_configurations(
             ADS_DATAFLOW_CONFIG_FILE_NAME,
@@ -668,6 +674,8 @@ def configure() -> None:
             ("bandwidth_mbps", ""),
             ("replica", ""),
             ("web_concurrency", ""),
+            ("memory_in_gbs", ""),
+            ("ocpus", "")
         ]
 
         _set_service_configurations(

ads/opctl/config/merger.py

@@ -59,6 +59,8 @@ def process(self, **kwargs) -> None:
         # 3. fill in values from default files under ~/.ads_ops
         self._fill_config_with_defaults(ads_config_path)
 
+        self._config_flex_shape_details()
+
         logger.debug(f"Config: {self.config}")
         return self
 
@@ -196,3 +198,47 @@ def _get_service_config(self, oci_profile: str, ads_config_folder: str) -> Dict:
                 f"{os.path.join(ads_config_folder, config_file)} does not exist. No config loaded."
             )
         return {}
+
+    def _config_flex_shape_details(self):
+        infrastructure = self.config["infrastructure"]
+        backend = self.config["execution"].get("backend", None)
+        if backend == BACKEND_NAME.JOB.value or backend == BACKEND_NAME.MODEL_DEPLOYMENT.value:
+            shape_name = infrastructure.get("shape_name", "")
+            if shape_name.endswith(".Flex"):
+                if (
+                    "ocpus" not in infrastructure or 
+                    "memory_in_gbs" not in infrastructure
+                ):
+                    raise ValueError(
+                        "Parameters `ocpus` and `memory_in_gbs` must be provided for using flex shape. "
+                        "Call `ads opctl config` to specify."
+                    )
+                infrastructure["shape_config_details"] = {
+                    "ocpus": infrastructure.pop("ocpus"),
+                    "memory_in_gbs": infrastructure.pop("memory_in_gbs")
+                }
+        elif backend == BACKEND_NAME.DATAFLOW.value:
+            executor_shape = infrastructure.get("executor_shape", "")
+            driver_shape = infrastructure.get("driver_shape", "")
+            data_flow_shape_config_details = [
+                "driver_shape_memory_in_gbs",
+                "driver_shape_ocpus",
+                "executor_shape_memory_in_gbs",
+                "executor_shape_ocpus"
+            ]
+            # executor_shape and driver_shape must be the same shape family
+            if executor_shape.endswith(".Flex") or driver_shape.endswith(".Flex"):
+                for parameter in data_flow_shape_config_details:
+                    if parameter not in infrastructure:
+                        raise ValueError(
+                            f"Parameters {parameter} must be provided for using flex shape. "
+                            "Call `ads opctl config` to specify."
+                        )
+                infrastructure["driver_shape_config"] = {
+                    "ocpus": infrastructure.pop("driver_shape_ocpus"),
+                    "memory_in_gbs": infrastructure.pop("driver_shape_memory_in_gbs")
+                }
+                infrastructure["executor_shape_config"] = {
+                    "ocpus": infrastructure.pop("executor_shape_ocpus"),
+                    "memory_in_gbs": infrastructure.pop("executor_shape_memory_in_gbs")
+                }

tests/unitary/with_extras/opctl/test_files/modeldeployment_conda.yaml

@@ -18,7 +18,10 @@ spec:
       logId: ocid1.log.oc1.iad.<unique_id>
       projectId: ocid1.datascienceproject.oc1.<unique_id>
       replica: 1
-      shapeName: VM.Standard.E2.4
+      shapeName: VM.Standard.E4.Flex
+      shapeConfigDetails:
+        ocpus: 1
+        memoryInGBs: 16
       webConcurrency: 10
     type: datascienceModelDeployment
   runtime:

tests/unitary/with_extras/opctl/test_files/modeldeployment_container.yaml

@@ -18,7 +18,10 @@ spec:
       logId: ocid1.log.oc1.iad.<unique_id>
       projectId: ocid1.datascienceproject.oc1.<unique_id>
       replica: 1
-      shapeName: VM.Standard.E2.4
+      shapeName: VM.Standard.E4.Flex
+      shapeConfigDetails:
+        ocpus: 1
+        memoryInGBs: 16
       webConcurrency: 10
     type: datascienceModelDeployment
   runtime:

tests/unitary/with_extras/opctl/test_opctl_cmds.py

@@ -45,11 +45,12 @@ def test_configure(self, confirm, prompt, monkeypatch):
                 ]
                 + ["abc"] * 8
                 + ["oci://bucket@namespace/path"]
+                + ["abc"] * 2
                 + ["abc"] * 4
                 + ["oci://bucket@namespace/path"]
-                + ["abc"] * 3
+                + ["abc"] * 7
                 + ["abc"] * 4
-                + ["abc"] * 8
+                + ["abc"] * 10
                 + ["1"]
                 + ["3"]
             )
@@ -115,11 +116,12 @@ def test_configure_in_notebook_session(self, confirm, prompt, monkeypatch):
                     ]
                     + ["abc"] * 8
                     + ["oci://bucket@namespace/path"]
+                    + ["abc"] * 2
                     + ["abc"] * 4
                     + ["oci://bucket@namespace/path"]
-                    + ["abc"] * 3
+                    + ["abc"] * 7
                     + ["abc"] * 4
-                    + ["abc"] * 8
+                    + ["abc"] * 10
                     + ["1"]
                     + ["3"]
                 )

tests/unitary/with_extras/opctl/test_opctl_config.py

@@ -169,6 +169,135 @@ def test_fill_config_with_defaults(self):
                 },
             }
 
+    def test_config_flex_shape_details(self):
+        config_one = {
+            "execution": {
+                "backend": "job",
+                "auth": "api_key",
+                "oci_config": DEFAULT_OCI_CONFIG_FILE,
+                "oci_profile": "PROFILE",
+                "conda_pack_folder": "~/condapack",
+                "conda_pack_os_prefix": "oci://bucket@namespace/path2",
+            },
+            "infrastructure": {
+                "compartment_id": "oci.compartmentid.abcd",
+                "project_id": "oci.projectid.abcd",
+                "shape_name": "VM.Standard.E2.4"
+            },
+        }
+
+        m = ConfigMerger(config_one)
+        m._config_flex_shape_details()
+
+        assert m.config == {
+            "execution": {
+                "backend": "job",
+                "auth": "api_key",
+                "oci_config": DEFAULT_OCI_CONFIG_FILE,
+                "oci_profile": "PROFILE",
+                "conda_pack_folder": "~/condapack",
+                "conda_pack_os_prefix": "oci://bucket@namespace/path2",
+            },
+            "infrastructure": {
+                "compartment_id": "oci.compartmentid.abcd",
+                "project_id": "oci.projectid.abcd",
+                "shape_name": "VM.Standard.E2.4"
+            },
+        }
+        
+        config_one["infrastructure"]["shape_name"] = "VM.Standard.E3.Flex"
+        m = ConfigMerger(config_one)
+
+        with pytest.raises(
+            ValueError, 
+            match="Parameters `ocpus` and `memory_in_gbs` must be provided for using flex shape. "
+                    "Call `ads opctl config` to specify."
+        ):
+            m._config_flex_shape_details()    
+
+        config_one["infrastructure"]["ocpus"] = 2
+        config_one["infrastructure"]["memory_in_gbs"] = 24
+        m = ConfigMerger(config_one)
+        m._config_flex_shape_details()
+
+        assert m.config == {
+            "execution": {
+                "backend": "job",
+                "auth": "api_key",
+                "oci_config": DEFAULT_OCI_CONFIG_FILE,
+                "oci_profile": "PROFILE",
+                "conda_pack_folder": "~/condapack",
+                "conda_pack_os_prefix": "oci://bucket@namespace/path2",
+            },
+            "infrastructure": {
+                "compartment_id": "oci.compartmentid.abcd",
+                "project_id": "oci.projectid.abcd",
+                "shape_name": "VM.Standard.E3.Flex",
+                "shape_config_details": {
+                    "ocpus": 2,
+                    "memory_in_gbs": 24
+                }
+            },
+        }
+
+        config_two = {
+            "execution": {
+                "backend": "dataflow",
+                "auth": "api_key",
+                "oci_config": DEFAULT_OCI_CONFIG_FILE,
+                "oci_profile": "PROFILE",
+                "conda_pack_folder": "~/condapack",
+                "conda_pack_os_prefix": "oci://bucket@namespace/path2",
+            },
+            "infrastructure": {
+                "compartment_id": "oci.compartmentid.abcd",
+                "project_id": "oci.projectid.abcd",
+                "executor_shape": "VM.Standard.E3.Flex",
+                "driver_shape": "VM.Standard.E3.Flex"
+            },
+        }
+
+        m = ConfigMerger(config_two)
+
+        with pytest.raises(
+            ValueError, 
+            match="Parameters driver_shape_memory_in_gbs must be provided for using flex shape. "
+                    "Call `ads opctl config` to specify."
+        ):
+            m._config_flex_shape_details()
+
+
+        config_two["infrastructure"]["driver_shape_memory_in_gbs"] = 36
+        config_two["infrastructure"]["driver_shape_ocpus"] = 4
+        config_two["infrastructure"]["executor_shape_memory_in_gbs"] = 48
+        config_two["infrastructure"]["executor_shape_ocpus"] = 5
+
+        m = ConfigMerger(config_two)
+        m._config_flex_shape_details()
+        assert m.config == {
+            "execution": {
+                "backend": "dataflow",
+                "auth": "api_key",
+                "oci_config": DEFAULT_OCI_CONFIG_FILE,
+                "oci_profile": "PROFILE",
+                "conda_pack_folder": "~/condapack",
+                "conda_pack_os_prefix": "oci://bucket@namespace/path2",
+            },
+            "infrastructure": {
+                "compartment_id": "oci.compartmentid.abcd",
+                "project_id": "oci.projectid.abcd",
+                "executor_shape": "VM.Standard.E3.Flex",
+                "executor_shape_config": {
+                    "ocpus": 5,
+                    "memory_in_gbs": 48
+                },
+                "driver_shape": "VM.Standard.E3.Flex",
+                "driver_shape_config": {
+                    "ocpus": 4,
+                    "memory_in_gbs": 36
+                }
+            },
+        }
 
 class TestConfigResolver:
     def test_resolve_operator_name(self):

tests/unitary/with_extras/opctl/test_opctl_model_deployment_backend.py

@@ -46,7 +46,7 @@ def config(self):
                 "project_id": "ocid1.datascienceproject.oc1.<unique_id>",
                 "log_group_id": "ocid1.loggroup.oc1.iad.<unique_id>",
                 "log_id": "ocid1.log.oc1.iad.<unique_id>",
-                "shape_name": "VM.Standard.E2.4",
+                "shape_name": "VM.Standard.E4.Flex",
                 "bandwidth_mbps": 10,
                 "replica": 1,
                 "web_concurrency": 10,