initial testable commit

Author: zeryx

adk/ADK.py

@@ -3,6 +3,10 @@
 import os
 import sys
 import Algorithmia
+import yaml
+import os
+import subprocess
+
 from adk.io import create_exception, format_data, format_response
 from adk.modeldata import ModelData
 
@@ -92,7 +96,24 @@ def process_local(self, local_payload, pprint):
         result = self.apply(local_payload)
         self.write_to_pipe(result, pprint=pprint)
 
-    def init(self, local_payload=None, pprint=print):
+    def mlops_initialize(self):
+        os.environ["MLOPS_SPOOLER_TYPE"] = "FILESYSTEM"
+        os.environ["MLOPS_FILESYSTEM_DIRECTORY"] = self.mlops_spool_dir
+        with open(f'{agents_dir}/conf/mlops.agent.conf.yaml') as f:
+            documents = yaml.load(f, Loader=yaml.FullLoader)
+        documents['mlopsUrl'] = DATAROBOT_ENDPOINT
+        documents['apiToken'] = DATAROBOT_API_TOKEN
+        with open(f'{agents_dir}/conf/mlops.agent.conf.yaml', 'w') as f:
+            yaml.dump(documents, f)
+        subprocess.call(f'{agents_dir}/bin/start-agent.sh')
+        check = subprocess.Popen([f'{agents_dir}/bin/status-agent.sh'], stdout=subprocess.PIPE)
+        check.terminate()
+
+
+
+    def init(self, local_payload=None, pprint=print, mlops=False):
+        if mlops and not self.is_local:
+            self.mlops_initialize()
         self.load()
         if self.is_local and local_payload is not None:
             if self.loading_exception:

adk/mlops.py

@@ -0,0 +1,36 @@
+import yaml
+import os
+import subprocess
+
+
+class MLOps(Object):
+    def __init__(self, endpoint, api_token, model_id, deployment_id):
+        self.token = api_token
+        self.endpoint = endpoint
+        self.model_id = model_id
+        self.deployment_id = deployment_id
+        self.spool_dir = "/tmp/ta"
+        self.agent_dir = "/opt/mlops-agent/datarobot_mlops_package-8.1.2"
+
+    def init(self):
+        with open(f'{self.agent_dir}/conf/mlops.agent.conf.yaml') as f:
+            documents = yaml.load(f, Loader=yaml.FullLoader)
+        documents['mlopsUrl'] = self.endpoint
+        documents['apiToken'] = self.token
+        with open(f'{agents_dir}/conf/mlops.agent.conf.yaml', 'w') as f:
+            yaml.dump(documents, f)
+
+        subprocess.call(f'{agents_dir}/bin/start-agent.sh')
+        check = subprocess.Popen([f'{agents_dir}/bin/status-agent.sh'], stdout=subprocess.PIPE)
+        output = check.stdout.readlines()
+        check.terminate()
+        if "DataRobot MLOps-Agent is running as a service." in output:
+            return True
+        else:
+            return False
+
+    def env_vars(self):
+        os.environ['MLOPS_DEPLOYMENT_ID'] = self.deployment_id
+        os.environ['MLOPS_MODEL_ID'] = self.model_id
+        os.environ['MLOPS_SPOOLER_TYPE'] = "FILESYSTEM"
+        os.environ['MLOPS_FILESYSTEM_DIRECTORY'] = "/tmp/ta"

adk/modeldata.py

@@ -2,17 +2,19 @@
 import json
 import hashlib
 from adk.classes import FileData
+from adk.mlops import MLOps
 
 
 class ModelData(object):
-    def __init__(self, client, model_manifest_path):
+    def __init__(self, client, model_manifest_path, mlops=False):
         self.manifest_reg_path = model_manifest_path
         self.manifest_frozen_path = "{}.freeze".format(self.manifest_reg_path)
         self.manifest_data = self.get_manifest()
         self.client = client
         self.models = {}
         self.usr_key = "__user__"
         self.using_frozen = True
+        self.use_mlops = mlops
 
     def __getitem__(self, key):
         return getattr(self, self.usr_key + key)
@@ -38,6 +40,8 @@ def available(self):
     def initialize(self):
         if self.client is None:
             raise Exception("Client was not defined, please define a Client when using Model Manifests.")
+        if self.use_mlops:
+            self.mlops_init()
         for required_file in self.manifest_data['required_files']:
             name = required_file['name']
             source_uri = required_file['source_uri']
@@ -88,16 +92,16 @@ def find_optional_model(self, file_name):
         else:
             self.models[file_name] = FileData(real_hash, local_data_path)
 
-
     def get_manifest(self):
         if os.path.exists(self.manifest_frozen_path):
             with open(self.manifest_frozen_path) as f:
                 manifest_data = json.load(f)
             if check_lock(manifest_data):
                 return manifest_data
             else:
-                raise Exception("Manifest FreezeFile Tamper Detected; please use the CLI and 'algo freeze' to rebuild your "
-                                "algorithm's freeze file.")
+                raise Exception(
+                    "Manifest FreezeFile Tamper Detected; please use the CLI and 'algo freeze' to rebuild your "
+                    "algorithm's freeze file.")
         elif os.path.exists(self.manifest_reg_path):
             with open(self.manifest_reg_path) as f:
                 manifest_data = json.load(f)
@@ -106,6 +110,18 @@ def get_manifest(self):
         else:
             return None
 
+    def mlops_init(self):
+        mlops = self.manifest_data['mlops']
+        model_id = mlops['model_id']
+        deployment_id = mlops['deployment_id']
+        datarobot_api_endpoint = mlops['datarobot_api_endpoint']
+
+        api_token = os.environ.get('DATAROBOT_MLOPS_API_TOKEN')
+        if api_token is None:
+            raise Exception("'DATAROBOT_MLOPS_API_TOKEN' environment variable not found.\nPlease ensure that you have a"
+                            "valid API token and add it as a secret to this algorithm.")
+        self.mlops = MLOps(datarobot_api_endpoint, api_token, model_id, deployment_id)
+
 
 def check_lock(manifest_data):
     expected_lock_checksum = manifest_data.get('lock_checksum')

tests/manifests/mlops_model_manifest.json

@@ -0,0 +1,9 @@
+{
+  "mlops": {
+    "model_id": "",
+    "deployment_id": "",
+    "datarobot_api_endpoint": "https://app.datarobot.com"
+  },
+  "required_models": [],
+  "optional_models": []
+}