functional, feature filled commit

Author: zeryx

adk/ADK.py

@@ -9,6 +9,7 @@
 
 from adk.io import create_exception, format_data, format_response
 from adk.modeldata import ModelData
+from adk.mlops import MLOps
 
 
 class ADK(object):
@@ -21,6 +22,7 @@ def __init__(self, apply_func, load_func=None, client=None):
         :param client: A Algorithmia Client instance that might be user defined,
          and is used for interacting with a model manifest file; if defined.
         """
+        self.mlops = None
         self.FIFO_PATH = "/tmp/algoout"
 
         if client:
@@ -43,10 +45,8 @@ def __init__(self, apply_func, load_func=None, client=None):
         self.load_result = None
         self.loading_exception = None
         self.manifest_path = "model_manifest.json"
-        self.model_data = self.init_manifest(self.manifest_path)
-
-    def init_manifest(self, path):
-        return ModelData(self.client, path)
+        self.mlops_path = "mlops.json"
+        self.model_data = ModelData(self.client, self.manifest_path)
 
     def load(self):
         try:
@@ -95,25 +95,18 @@ def write_to_pipe(self, payload, pprint=print):
     def process_local(self, local_payload, pprint):
         result = self.apply(local_payload)
         self.write_to_pipe(result, pprint=pprint)
-
-    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 mlops_init(self):
+        mlops_token = os.environ.get("DATAROBOT_MLOPS_API_TOKEN", None)
+        if mlops_token:
+            self.mlops = MLOps(mlops_token, self.mlops_path)
+            self.mlops.init()
+        else:
+            raise Exception("'DATAROBOT_MLOPS_API_TOKEN' was not found, please set to use mlops.")
 
     def init(self, local_payload=None, pprint=print, mlops=False):
         if mlops and not self.is_local:
-            self.mlops_initialize()
+            self.mlops_init()
         self.load()
         if self.is_local and local_payload is not None:
             if self.loading_exception:

adk/mlops.py

@@ -1,36 +1,44 @@
 import yaml
+import json
 import os
 import subprocess
 
 
-class MLOps(Object):
-    def __init__(self, endpoint, api_token, model_id, deployment_id):
+class MLOps(object):
+    spool_dir = "/tmp/ta"
+    agent_dir = "/opt/mlops-agent/datarobot_mlops_package-8.1.2"
+
+    def __init__(self, api_token, path):
         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"
+        if os.path.exists(path):
+            with open(path) as f:
+                mlops_config = json.load(f)
+        else:
+            raise Exception("'mlops.json' file does not exist, but mlops was requested.")
+        if not os.path.exists(agent_dir):
+            raise Exception("environment is not configured for mlops.\nPlease select a valid mlops enabled environment.")
+        self.endpoint = mlops_config['datarobot_api_endpoint']
+        self.model_id = mlops_config['model_id']
+        self.deployment_id = mlops_config['deployment_id']
 
     def init(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"
+
         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:
+        with open(f'{self.agent_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()
+        subprocess.call(f'{self.agent_dir}/bin/start-agent.sh')
+        check = subprocess.Popen([f'{self.agent_dir}/bin/status-agent.sh'], stdout=subprocess.PIPE)
+        output = check.stdout.readlines()[0]
         check.terminate()
-        if "DataRobot MLOps-Agent is running as a service." in output:
+        if b"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"
+            raise Exception(output)

adk/modeldata.py

@@ -2,19 +2,17 @@
 import json
 import hashlib
 from adk.classes import FileData
-from adk.mlops import MLOps
 
 
 class ModelData(object):
-    def __init__(self, client, model_manifest_path, mlops=False):
+    def __init__(self, client, model_manifest_path):
         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)
@@ -40,8 +38,6 @@ 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']
@@ -110,18 +106,6 @@ 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')