configured for use with customizable mlflow model

label_studio_ml/api.py

@@ -1,6 +1,12 @@
 import hmac
+import json
 import logging
 import os
+import dagshub
+import mlflow
+import base64
+import cloudpickle
+from dagshub.data_engine import datasources
 
 from flask import Flask, request, jsonify, Response
 
@@ -11,25 +17,34 @@
 logger = logging.getLogger(__name__)
 
 _server = Flask(__name__)
-MODEL_CLASS = LabelStudioMLBase
 BASIC_AUTH = None
 
 
-def init_app(model_class, basic_auth_user=None, basic_auth_pass=None):
- global MODEL_CLASS
+def init_app(model_instance, basic_auth_user=None, basic_auth_pass=None):
+ global model
  global BASIC_AUTH
 
- if not issubclass(model_class, LabelStudioMLBase):
- raise ValueError('Inference class should be the subclass of ' + LabelStudioMLBase.__class__.__name__)
-
- MODEL_CLASS = model_class
+ model = model_instance
  basic_auth_user = basic_auth_user or os.environ.get('BASIC_AUTH_USER')
  basic_auth_pass = basic_auth_pass or os.environ.get('BASIC_AUTH_PASS')
  if basic_auth_user and basic_auth_pass:
  BASIC_AUTH = (basic_auth_user, basic_auth_pass)
 
  return _server
 
+@_server.post('/configure')
+@exception_handler
+def _configure():
+ args = json.loads(request.get_json())
+ dagshub.init(args['repo'], args['username']) # user-level privileged auth token
+ ls_model = mlflow.pyfunc.load_model(f'models:/{args["model"]}/{args["version"]}')
+
+ model.configure(ls_model, *[cloudpickle.loads(base64.b64decode(args[hook])) for hook in ['pre_hook', 'post_hook']])
+ # model.api = dagshub.common.api.repo.RepoAPI(f'https://dagshub.com/{args["username"]}/{args["repo"]}', host=args['host'])
+
+ model.ds = datasources.get_datasource(args['datasource_repo'], args['datasource_name'])
+ model.dp_map = model.ds.all().dataframe[['path', 'datapoint_id']]
+ return []
 
 @_server.route('/predict', methods=['POST'])
 @exception_handler
@@ -61,8 +76,8 @@ def _predict():
  params = data.get('params', {})
  context = params.pop('context', {})
 
- model = MODEL_CLASS(project_id=project_id,
-  label_config=label_config)
+ model.project_id = project_id
+ model.use_label_config(label_config)
 
  # model.use_label_config(label_config)
 
@@ -96,8 +111,8 @@ def _setup():
  project_id = data.get('project').split('.', 1)[0]
  label_config = data.get('schema')
  extra_params = data.get('extra_params')
- model = MODEL_CLASS(project_id=project_id,
-  label_config=label_config)
+ model.project_id = project_id
+ model.use_label_config(label_config)
 
  if extra_params:
  model.set_extra_params(extra_params)
@@ -122,7 +137,8 @@ def webhook():
  return jsonify({'status': 'Unknown event'}), 200
  project_id = str(data['project']['id'])
  label_config = data['project']['label_config']
- model = MODEL_CLASS(project_id, label_config=label_config)
+ model.project_id = project_id
+ model.use_label_config(label_config)
  model.fit(event, data)
  return jsonify({}), 201
 
@@ -133,7 +149,6 @@ def webhook():
 def health():
  return jsonify({
  'status': 'UP',
- 'model_class': MODEL_CLASS.__name__
  })