Merge branch 'main' into watson_nlp

IBM · Nov 10, 2023 · 0b2b3ee · 0b2b3ee
2 parents 334d674 + bd03b83
commit 0b2b3ee
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 1 deletion.
diff --git a/mlflow_watsonml/deploy.py b/mlflow_watsonml/deploy.py
@@ -142,6 +142,7 @@ def create_deployment(
             - "conda_yaml" : filepath of conda.yaml file
             - "custom_packages": a list of str - zip file paths of the packages
             - "rewrite_software_spec": bool whether to rewrite the software spec
+            - "hardware_spec_name" : name of the hardware specification to use (Default: XS)
         endpoint : str
             deployment space name
 
@@ -206,13 +207,27 @@ def create_deployment(
 
         batch = config.get("batch", False)
 
+        hardware_spec_name = config.get("hardware_spec_name")
+        if hardware_spec_name is not None:
+            hardware_spec_id = client.hardware_specifications.get_id_by_name(
+                hardware_spec_name
+            )
+            if hardware_spec_id == "Not Found":
+                LOGGER.warn(
+                    f"Hardware Specification - {hardware_spec_name} not found. Using default."
+                )
+                hardware_spec_id = None
+        else:
+            hardware_spec_id = None
+
         deployment_details = deploy(
             client=client,
             name=name,
             artifact_id=artifact_id,
             revision_id=revision_id,
             batch=batch,
             environment_variables=environment_variables,
+            hardware_spec_id=hardware_spec_id,
         )
 
         return deployment_details

diff --git a/mlflow_watsonml/wml.py b/mlflow_watsonml/wml.py
@@ -18,6 +18,7 @@ def deploy(
     revision_id: str,
     batch: bool = False,
     environment_variables: Optional[Dict] = None,
+    hardware_spec_id: Optional[str] = None,
 ) -> Dict:
     """Create a new WML deployment
 
@@ -46,7 +47,11 @@ def deploy(
             client.deployments.ConfigurationMetaNames.NAME: name,
             client.deployments.ConfigurationMetaNames.BATCH: {},
             client.deployments.ConfigurationMetaNames.CUSTOM: environment_variables,
-            client.deployments.ConfigurationMetaNames.HARDWARE_SPEC: {},
+            client.deployments.ConfigurationMetaNames.HARDWARE_SPEC: {
+                "id": hardware_spec_id
+            }
+            if hardware_spec_id is not None
+            else {},
             client.deployments.ConfigurationMetaNames.ASSET: {
                 "id": artifact_id,
                 "rev": revision_id,
@@ -56,6 +61,11 @@ def deploy(
         deployment_props = {
             client.deployments.ConfigurationMetaNames.NAME: name,
             client.deployments.ConfigurationMetaNames.CUSTOM: environment_variables,
+            client.deployments.ConfigurationMetaNames.HARDWARE_SPEC: {
+                "id": hardware_spec_id
+            }
+            if hardware_spec_id is not None
+            else {},
             client.deployments.ConfigurationMetaNames.ONLINE: {},
             client.deployments.ConfigurationMetaNames.ASSET: {
                 "id": artifact_id,