Merge pull request #91 from mindsdb/byom-support

Support byom ml engine upload

Author: ea-rus

mindsdb_sdk/connectors/rest_api.py

@@ -141,6 +141,19 @@ def upload_file(self, name: str, df: pd.DataFrame):
         )
         _raise_for_status(r)
 
+    @_try_relogin
+    def upload_byom(self, name: str, code: str, requirements: str):
+
+        url = self.url + f'/api/handlers/byom/{name}'
+        r = self.session.put(
+            url,
+            files={
+                'code': code,
+                'modules': requirements,
+            }
+        )
+        _raise_for_status(r)
+
     def status(self) -> dict:
 
         r = self.session.get(self.url + f'/api/status')

mindsdb_sdk/ml_engines.py

@@ -45,6 +45,16 @@ class MLEngines(CollectionBase):
 
     >>>  con.ml_engines.drop('openai1')
 
+    Upload BYOM model. After uploading a new ml engin will be availbe to create new model from it.
+
+    >>> model_code = open('/path/to/model/code').read()
+    >>> model_requirements = open('/path/to/model/requirements').read()
+    >>> ml_engine = con.ml_engines.create_byom(
+    ...    'my_byom_engine',
+    ...    code=model_code,
+    ...    requirements=model_requirements
+    ...)
+
     """
 
     def __init__(self, api):
@@ -101,6 +111,25 @@ def create(self, name: str, handler: Union[str, Handler], connection_data: dict
 
         return MLEngine(name, handler, connection_data)
 
+    def create_byom(self, name: str, code: str, requirements: Union[str, List[str]] = None):
+        """
+        Create new BYOM ML engine and return it
+
+        :param code: model python code in string
+        :param requirements: requirements for model. Optional if there is no special requirements.
+           It can be content of 'requirement.txt' file or list of strings (item for every requirement).
+        :return: created BYOM ml engine object
+        """
+
+        if requirements is None:
+            requirements = ''
+        elif isinstance(requirements, list):
+            requirements = '\n'.join(requirements)
+
+        self.api.upload_byom(name, code, requirements)
+
+        return MLEngine(name, 'byom', {})
+
     def drop(self, name: str):
         """
         Drop ml engine by name

tests/test_sdk.py

@@ -700,6 +700,27 @@ def test_flow(self, mock_post, mock_put):
         con.ml_engines.drop('openai1')
         check_sql_call(mock_post, 'DROP ML_ENGINE openai1')
 
+        # byom
+        model = '''
+import pandas as pd
+
+class CustomPredictor():
+
+    def train(self, df, target_col, args=None):
+        
+        self.target_col=target_col 
+        
+     def predict(self, df):
+        return pd.Dataframe([{'predict': self.target_col}])
+'''
+        requirements = '''pandas'''
+
+        con.ml_engines.create_byom('b1', model, requirements)
+        call_args = mock_put.call_args
+        assert call_args[0][0] == 'https://cloud.mindsdb.com/api/handlers/byom/b1'
+        assert call_args[1]['files']['code'] == model
+        assert call_args[1]['files']['modules'] == requirements
+
     def check_project(self, project, database):
         self.check_project_views( project, database)