A Simple Way to Manage Data Science Assets with IBM Watson Machine Learning API

wml_utils.py

import requests
# from watson_machine_learning_client import WatsonMachineLearningAPIClient
from ibm_watson_machine_learning import APIClient

class WMLUtils:
    def __init__(self, wml_credentials):
#         self.client = WatsonMachineLearningAPIClient(wml_credentials)
        self.client = APIClient(wml_credentials)
        
    def set_deployment_space(self, deployment_space):
        """
        Creates a deployment space if not exist and sets it as the default deployment space.
        
        Args:
          deployment_space: Deployment space name

        Returns:
          space_id: The space id that was created or found based on the deployment space name
        """
        deployment_space_dict = {space["entity"]["name"]:space["metadata"]["id"] for space in self.client.spaces.get_details()["resources"]}
        if deployment_space in deployment_space_dict:
            space_id = deployment_space_dict[deployment_space]
        else:
            metadata = {
                self.client.spaces.ConfigurationMetaNames.NAME: deployment_space,
            }
            space_details = self.client.spaces.store(meta_props=metadata)
            space_id = self.client.spaces.get_uid(space_details)
        
        print(f"Deployment space: {deployment_space} - {space_id}")
        self.client.set.default_space(space_id)
        return space_id
    
    def store_model(self, model_name, model, model_type, software_specifications, update_if_exist=True):
        """
        If update exist is set to true, updates the model based on the model name specified. 
        If update exist is set to false or the model name does not exist, creates a new model.
        
        Args:
          model_name: Model name to be stored
          model: Model to be stored
          model_type: Based on client.repository.ModelMetaNames.show()
          software_specifications: Based on client.software_specifications.list()

        Returns:
          model_id: The model id that was created or updated based on the model name
        """
        model_dict = {model["metadata"]["name"]:model["metadata"]["id"] for model in self.client.repository.get_model_details()["resources"]}
        
        if model_name in model_dict and update_if_exist:
            metadata = {
                self.client.repository.ModelMetaNames.NAME: model_name,
            }
            model_details = self.client.repository.update_model(model_dict[model_name], updated_meta_props=metadata, update_model=model)
        else:
            metadata = {
                self.client.repository.ModelMetaNames.NAME: model_name,
                self.client.repository.ModelMetaNames.TYPE: model_type,
                self.client.repository.ModelMetaNames.SOFTWARE_SPEC_UID: self.client.software_specifications.get_uid_by_name(software_specifications)
            }
            model_details = self.client.repository.store_model(model, meta_props=metadata)
        
        model_id = self.client.repository.get_model_uid(model_details)
        print(f"Model: {model_name} - {model_id}")
        return model_id
    
    def deploy_model(self, deployment_name, deployment_type, model_id, update_if_exist=True):
        """
        If update exist is set to true, updates the deployment based on the deployment name specified. 
        If update exist is set to false or the deployment name does not exist, creates a new deployment.
        
        Args:
          deployment_name: Deployment name to be stored
          deployment_type: Deployment type
          model_id: Model id to be deployed
          
        Returns:
          deployment_id: The deployment_id id that was created or updated based on the deployment name
        """
        deployment_dict = {deployment["metadata"]["name"]:deployment["metadata"]["id"] for deployment in self.client.deployments.get_details()["resources"]}
        
        if deployment_name in deployment_dict and update_if_exist:
            metadata = {
                self.client.deployments.ConfigurationMetaNames.ASSET: {"id": model_id}
            }
            model_deployment_details = self.client.deployments.update(deployment_dict[deployment_name], changes=metadata)
        else:
            meta_props = {
                self.client.deployments.ConfigurationMetaNames.NAME: deployment_name,
                self.client.deployments.ConfigurationMetaNames.ONLINE: {}
            }
            model_deployment_details = self.client.deployments.create(model_id, meta_props)
        deployment_id = self.client.deployments.get_uid(model_deployment_details)
        print(f"Deployment: {deployment_name} - {deployment_id}")
        return deployment_id
    
    
    def get_deployment_id(self, function_name):
        deployment_dict = {deployment["metadata"]["name"]:deployment["metadata"]["id"] for deployment in self.client.deployments.get_details()["resources"]}
        return deployment_dict[function_name]
    
    
    def upload_asset(self, assets):
        data = { "query": "*:*" }
        response = requests.post(self.client.data_assets._href_definitions.get_search_asset_href(), params=self.client.data_assets._client._params(), headers=self.client.data_assets._client._get_headers(), json=data)
        self.client.data_assets._handle_response(200, "list assets", response)
        asset_details = self.client.data_assets._handle_response(200, "list assets", response)["results"]
        assets_dict = {asset_detail["metadata"]["name"]:asset_detail["metadata"]["asset_id"] for asset_detail in asset_details}
        for name, value in assets.items():
            file_path = value["file_path"]        
            if name in assets_dict:
                asset_id = assets_dict[name]
            else:
                asset_details = self.client.data_assets.create(name=name, file_path=file_path)
                asset_id = self.client.data_assets.get_uid(asset_details)
            value["asset_id"] = asset_id
        return assets