Added flex shape for opctl config

Author: lu-ohai

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_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):