def test_create_model():
with requests_mock.mock() as m:
rh = MlOpsRestFactory().get_rest_helper(MLOpsMode.AGENT)
model_id = "model_5906255e-0a3d-4fef-8653-8d41911264fb"
m.get(rh.url_get_uuid("model"), json={"id": model_id})
ion = ION()
ion.id = "bdc2ee10-767c-4524-ba72-8268a3894bff"
mh = ModelHelper(rest_helper=rh, ion=ion, stats_helper=None)
model_data = "MODEL_DATA"
model = mh.create_model(name="my model", model_format=ModelFormat.TEXT, description="test model")
model_file = os.path.join(os.path.sep, "tmp", str(uuid.uuid4()))
f = open(model_file, 'w')
f.write(model_data)
f.close()
model.set_model_path(model_file)
assert model.get_id() == model_id
os.remove(model_file)
rh.done()
def test_publish_model_rest():
with requests_mock.mock() as m:
rh = MlOpsRestFactory().get_rest_helper(MLOpsMode.AGENT, mlops_server="localhost", mlops_port="4567")
model_id = "model_5906255e-0a3d-4fef-8653-8d41911264fb"
m.post('http://localhost:4567/models', json=model_id)
m.get(rh.url_get_uuid("model"), json={"id": model_id})
ion = ION()
ion.id = "bdc2ee10-767c-4524-ba72-8268a3894bff"
mh = ModelHelper(rest_helper=rh, ion=ion, stats_helper=None)
model_data = "MODEL_DATA"
model = mh.create_model(name="my model", model_format=ModelFormat.TEXT, description="test model",
user_defined="whatever I want goes here")
model_file = os.path.join(os.path.sep, "tmp", str(uuid.uuid4()))
f = open(model_file, 'w')
f.write(model_data)
f.close()
model.set_model_path(model_file)
my_id = mh.publish_model(model, None)
os.remove(model_file)
assert (model_id == my_id)
rh.done()
def test_convert_models_json_dict_to_dataframe():
rh = MlOpsRestFactory().get_rest_helper(MLOpsMode.AGENT)
ion = ION()
ion.id = "bdc2ee10-767c-4524-ba72-8268a3894bff"
mh = ModelHelper(rest_helper=rh, ion=ion, stats_helper=None)
df = mh.convert_models_json_dict_to_dataframe(models_list_json_dict)
assert len(df) == 2
rh.done()
def test_publish_model():
expected_models_list_json_dict = [
{
models.json_fields.MODEL_ID_FIELD: '',
models.json_fields.MODEL_NAME_FIELD: 'my model name',
models.json_fields.MODEL_FORMAT_FIELD: 'Text',
models.json_fields.MODEL_VERSION_FIELD: '',
models.json_fields.MODEL_DESCRIPTION_FIELD: 'test model',
models.json_fields.MODEL_TRAIN_VERSION_FIELD: '',
models.json_fields.MODEL_SIZE_FIELD: 10,
models.json_fields.MODEL_OWNER_FIELD: '',
models.json_fields.MODEL_CREATED_ON_FIELD: None,
models.json_fields.MODEL_FLAG_VALUES_FIELD: [],
models.json_fields.MODEL_ANNOTATIONS_FIELD: {"custom_data": "my content"},
models.json_fields.MODEL_ACTIVE_FIELD: False
}
]
rh = MlOpsRestFactory().get_rest_helper(MLOpsMode.STAND_ALONE)
ion = ION()
mh = ModelHelper(rest_helper=rh, ion=ion, stats_helper=None)
model_data = "MODEL_DATA"
model = mh.create_model(name="my model name", model_format=ModelFormat.TEXT, description="test model")
model.set_annotations({"custom_data": "my content"})
model_file = os.path.join(os.path.sep, "tmp", str(uuid.uuid4()))
f = open(model_file, 'w')
f.write(model_data)
f.close()
model.set_model_path(model_file)
my_id = mh.publish_model(model, None)
os.remove(model_file)
assert my_id == model.get_id()
expected_models_list_json_dict[0][models.json_fields.MODEL_ID_FIELD] = my_id
ret_data = mh.download_model(my_id)
assert ret_data == model_data
result_model_list = mh.fetch_all_models_json_dict()
actual_json_dumps = json.dumps(result_model_list, sort_keys=True, indent=2)
local_json_dump = json.dumps(expected_models_list_json_dict, sort_keys=True, indent=2)
print("Expected_Dumps: {}".format(local_json_dump))
print("Actual_Dumps: {}".format(actual_json_dumps))
assert expected_models_list_json_dict == result_model_list
with pytest.raises(MLOpsException):
mh.publish_model("Not a model", None)
rh.done()
def test_model_list_dict_from_json():
with requests_mock.mock() as m:
m.get('http://localhost:3456/v1/models', json=models_list_json_dict)
rh = MlOpsRestFactory().get_rest_helper(MLOpsMode.AGENT)
ion = ION()
ion.id = "bdc2ee10-767c-4524-ba72-8268a3894bff"
mh = ModelHelper(rest_helper=rh, ion=ion, stats_helper=None)
result_model_list = mh.fetch_all_models_json_dict()
print("Type is: {}".format(type(result_model_list)))
print("result_model_list: {}".format(result_model_list))
json_str_orig = json.dumps(models_list_json_dict, sort_keys=True, indent=2)
json_str_got = json.dumps(result_model_list, sort_keys=True, indent=2)
assert json_str_orig == json_str_got
rh.done()
def test_get_models_with_filter():
with requests_mock.mock() as m:
m.get('http://localhost:3456/v1/models', json=models_list_json_dict)
rh = MlOpsRestFactory().get_rest_helper(MLOpsMode.AGENT)
ion = ION()
ion.id = "bdc2ee10-767c-4524-ba72-8268a3894bff"
mh = ModelHelper(rest_helper=rh, ion=ion, stats_helper=None)
mf = ModelFilter()
mf.time_window_start = datetime.utcfromtimestamp(1518460571573 / 1000)
mf.time_window_end = datetime.utcfromtimestamp(1518460577573 / 1000)
filtered_models = mh.get_models_dataframe(model_filter=mf, download=False)
assert len(filtered_models) == 1
print(filtered_models[[models.json_fields.MODEL_NAME_FIELD, models.json_fields.MODEL_CREATED_ON_FIELD]])
rh.done()
def test_all_alerts():
mlops_ctx = build_ion_ctx()
with requests_mock.mock() as m:
m.get('http://localhost:3456/events', json=alerts_list)
rest_helper = MlOpsRestFactory().get_rest_helper(MLOpsMode.AGENT)
event_helper = EventBroker(mlops_ctx, None)
ef = EventFilter()
alerts = event_helper.get_events(ef)
assert len(alerts) == 2
alert_df = alerts[alerts.id == AlertsInfo.ALERT_0_ID]
assert alert_df.iloc[0]["node"] == ION1.NODE_0_ID
rest_helper.done()
def test_get_models_with_filter_2():
with requests_mock.mock() as m:
m.get('http://localhost:3456/v1/models', json=models_list_json_dict)
rh = MlOpsRestFactory().get_rest_helper(MLOpsMode.AGENT)
ion = ION()
ion.id = "13445bb4-535a-4d45-b2f2-77293026e3da"
mh = ModelHelper(rest_helper=rh, ion=ion, stats_helper=None)
model_id_to_filter = '8c95deaf-87e4-4c21-bc92-e5b1a0454f9a'
mf = ModelFilter()
mf.id = model_id_to_filter
filtered_models = mh.get_models_dataframe(model_filter=mf, download=False)
print(filtered_models[[models.json_fields.MODEL_ID_FIELD, models.json_fields.MODEL_FORMAT_FIELD]])
assert len(filtered_models) == 1
assert filtered_models.iloc[0][models.json_fields.MODEL_FORMAT_FIELD] == 'TEXT'
assert filtered_models.iloc[0][models.json_fields.MODEL_ID_FIELD] == model_id_to_filter
rh.done()
def test_get_models_with_filter_3():
with requests_mock.mock() as m:
m.get('http://localhost:3456/models', json=models_list_json_dict)
rh = MlOpsRestFactory().get_rest_helper(MLOpsMode.AGENT)
ion = ION()
ion.id = "bdc2ee10-767c-4524-ba72-8268a3894bff"
mh = ModelHelper(rest_helper=rh, ion=ion, stats_helper=None)
mf = ModelFilter()
mf.time_window_start = datetime.utcfromtimestamp(1518460571573 / 1000)
mf.time_window_end = datetime.utcfromtimestamp(1518460577573 / 1000)
mf.pipeline_instance_id = ['94bf382b-47d5-4b80-b76c-3bca862e6e23', 'asdf']
filtered_models = mh.get_models_dataframe(model_filter=mf, download=False)
assert len(filtered_models) == 1
print(filtered_models[["name", "createdTimestamp", "pipelineInstanceId"]])
# No model found
mf.id = "111111111111111"
filtered_models = mh.get_models_dataframe(model_filter=mf, download=False)
assert len(filtered_models) == 0
rh.done()
def test_publish_model():
rh = MlOpsRestFactory().get_rest_helper(MLOpsMode.STAND_ALONE)
ion = ION()
ion.id = "bdc2ee10-767c-4524-ba72-8268a3894bff"
local_models_list_json_dict[0]["workflowRunId"] = ion.id
mh = ModelHelper(rest_helper=rh, ion=ion, stats_helper=None)
model_data = "MODEL_DATA"
model = mh.create_model(name="my model", model_format=ModelFormat.TEXT, description="test model",
user_defined="whatever I want goes here")
model_file = os.path.join(os.path.sep, "tmp", str(uuid.uuid4()))
f = open(model_file, 'w')
f.write(model_data)
f.close()
model.set_model_path(model_file)
my_id = mh.publish_model(model, None)
os.remove(model_file)
assert my_id == model.get_id()
local_models_list_json_dict[0]["modelId"] = my_id
ret_data = mh.download_model(my_id)
assert ret_data == model_data
result_model_list = mh.fetch_all_models_json_dict()
actual_json_dumps = json.dumps(result_model_list, sort_keys=True, indent=2)
local_json_dump = json.dumps(local_models_list_json_dict, sort_keys=True, indent=2)
print("Expected_Dumps: {}".format(local_json_dump))
print("Actual_Dumps: {}".format(actual_json_dumps))
assert local_models_list_json_dict == result_model_list
with pytest.raises(MLOpsException):
mh.publish_model("Not a model", None)
rh.done()
class MLOpsCtx(BaseObj):
"""
Provide context information for MLOps library.
This is an internal class which should not be exposed to users of MLOps.
The object contains the structure of the ION and groups and such.
"""
def __init__(self, config, mode=None):
"""
Perform initialization of the MLOpsCtx.
It expects configuration from the environment to arrive in the standard usage mode.
:param config: :class:`ConfigInfo` for this MLOps instantiation
:param mode: python or pyspark
:return:
: raises MLOpsException for invalid configurations
"""
super(MLOpsCtx, self).__init__(__name__)
self._info("MLOpsCtx __init__ called")
self._info("Config\n{}".format(config))
self._agent_list = []
self._rest_helper = None
self._ci = config
self._mode = mode
self._ion = None # Will contain an ION class corresponding to the active ION which this mlops is part of
self._ees_dict = {} # Will contain all ees defined with the agent inside
self._agents_dict = {}
self._rest_helper = MlOpsRestFactory().get_rest_helper(self._mode, self._ci.mlops_server, self._ci.mlops_port, self._ci.token)
if self._mode == MLOpsMode.AGENT or self._mode == MLOpsMode.REST_ACCUMULATOR:
# In agent mode, we talk with the agent and use the mlops prefix to the http requests
self._validate_config()
self._info("Agent mode")
self._rest_helper.set_prefix(Constants.URL_MLOPS_PREFIX)
json_dict = self._detect_ion_structure()
health_json_dict = self._fetch_health_thresholds()
self._detect_ees_and_agents()
self._build_ion_obj(json_dict)
self._build_health_obj(health_json_dict)
elif self._mode == MLOpsMode.ATTACH:
# In attach mode, we connect either to the ZK or to the server directly
if self._ci.zk_host:
self._detect_mlops_server_via_zk()
self._validate_config()
self._info("In pm mode - will try to connect to server")
ion_json_dict = self._detect_ion_structure()
health_json_dict = self._fetch_health_thresholds()
self._detect_ees_and_agents()
self._build_ion_obj(ion_json_dict)
self._build_health_obj(health_json_dict)
elif self._mode == MLOpsMode.STAND_ALONE:
# In stand alone mode, we do not have a valid ION structure
self._logger.info("In stand-alone mode: ctx data will not be available")
self._set_stand_alone_values()
else:
raise MLOpsException("Unsupported operation mode: {}".format(self._mode))
def _fetch_health_thresholds(self):
return self._rest_helper.get_health_thresholds(self._ci.ion_id)
def _set_stand_alone_values(self):
self._ion = ION()
self._ion.id = 1
self._ion.name = "ION_1"
def _validate_config(self):
"""
Validate that all config information is present
:return:
:raises MLOpsException for invalid configurations
"""
if self._ci.token is None:
raise MLOpsException("Internal Error: No auth token provided")
if self._ci.mlops_server is None or self._ci.mlops_port is None:
raise MLOpsException("MLOps server host or port were not provided")
if self._ci.ion_id is None:
MLOpsException("{} instance id not provided".format(Constants.ION_LITERAL))
def _detect_mlops_server_via_zk(self):
"""
Detect the active mlops server via the ZK
:return:
"""
zk = None
try:
zk = KazooClient(hosts=self._ci.zk_host, read_only=True)
zk.start()
if zk.exists(Constants.MLOPS_ZK_ACTIVE_HOST_PORT):
data, stat = zk.get(Constants.MLOPS_ZK_ACTIVE_HOST_PORT)
eco_host_port = data.decode("utf-8").split(':')
if len(eco_host_port) is 2:
self._ci.mlops_server = eco_host_port[0]
self._ci.mlops_port = eco_host_port[1]
else:
raise MLOpsException("Internal Error: Invalid zookeeper active server entry, host_port: {}"
.format(eco_host_port))
else:
raise MLOpsException("Unable to connect to the active MLOps server, zk_host: {}"
.format(self._ci.zk_host))
except Exception as e:
raise MLOpsException("{}, zk_host: {}".format(e, self._ci.zk_host))
finally:
if zk:
zk.stop()
@staticmethod
def _search_list_dict(kv, key, value):
for x in kv:
if x[key] == value:
return x
return None
def _detect_ion_structure(self):
"""
Detect the current ion structure (pipeline, groups, agents, etc.)
:return:
:raises MLOpsException if ION or other structures are not found
"""
self._info("Detecting {} structure".format(Constants.ION_LITERAL))
# This is the max number of retries to wait until the ION is running.
# The pipelineInstance part of the workflow description does not appear until the ION
# switches to RUNNING state. For this reason, the code loop until the pipelineInstances part
# appears in the JSON.
max_tries = Constants.WAIT_FOR_PIPELINE_INSTANCE_TO_APPEAR_TIMEOUT
wf_instance = {}
found = False
for idx in range(0, max_tries):
wf_instance = self._rest_helper.get_workflow_instance(self._ci.ion_id)
if wf_instance is None:
raise MLOpsException("Could not locate {} instance {}".format(
Constants.ION_LITERAL, self._ci.ion_id))
self._debug("{} status: {}".format(Constants.ION_LITERAL, wf_instance['status']))
if IONJsonConstants.PIPELINE_INSTANCES_SECTION in wf_instance:
found = True
break
self._info("Could not find {} in workflow json - try {}".format(
IONJsonConstants.PIPELINE_INSTANCES_SECTION, idx))
time.sleep(1)
if found is False:
raise MLOpsException("Could not find {} section in workflow information".format(
IONJsonConstants.PIPELINE_INSTANCES_SECTION))
ion_json_dict = wf_instance
self._debug("workflow: {}".format(ion_json_dict))
return ion_json_dict
def _detect_ees_and_agents(self):
ees_json_dict = self._rest_helper.get_ees()
agents_json_dict = self._rest_helper.get_agents()
self._debug("Agents JSON:\n{}\n\n".format(agents_json_dict))
self._debug("EEs JSON:\n{}\n\n".format(ees_json_dict))
# Generating a dict of all agents by ID
for agent_json in agents_json_dict:
agent_obj = Agent()
agent_obj.id = str(agent_json["id"])
agent_obj.hostname = str(agent_json["address"])
self._agents_dict[agent_obj.id] = agent_obj
for ee_json in ees_json_dict:
ee = EE()
ee.name = str(ee_json["name"])
ee.id = str(ee_json["id"])
ee.agent_id = str(ee_json["agentId"])
# get agent object we created above in the agent_dict
if ee.agent_id not in self._agents_dict:
raise MLOpsException("EE {} contains Agent {} which is not in global agent list".format(
ee.name, ee.agent_id))
agent_obj = self._agents_dict[ee.agent_id]
ee.agents.append(agent_obj)
ee.agent_by_id[ee.agent_id] = agent_obj
ee.agent_by_hostname[agent_obj.hostname] = agent_obj
self._ees_dict[ee.id] = ee
self._logger.info("EE:\n{}".format(ee))
def _build_ion_obj(self, ion_json_dict):
ion_builder = IONBuilder()
self._ion = ion_builder.build_from_dict(ion_json_dict)
# This info line is important so we can understand what was the ion structure if errors happens at customer site
self._info("{}:\n{}".format(Constants.ION_LITERAL, self._ion))
self._info("---------------------")
def _build_health_obj(self, health_json_dict):
self._ion.policy.set_thresholds(health_json_dict[IONJsonConstants.ION_GLOBAL_THRESHOLD_TAG], health_json_dict[IONJsonConstants.ION_CANARY_THRESHOLD_TAG])
def rest_helper(self):
return self._rest_helper
def ion(self):
"""
Return a copy of the ion object
:return: ION object
:rtype: ION
"""
return copy.deepcopy(self._ion)
def ion_id(self):
"""
Return the ION id
:return:
"""
return self._ion.id
def ion_node_id(self):
"""
Return the current node id this code is running in
:return: ION node id
"""
return self._ci.ion_node_id
def current_node(self):
if self._ci.ion_node_id not in self._ion.node_by_id:
raise MLOpsException("Current node id: [{}] is not detected in {}".format(
self._ci.ion_node_id, Constants.ION_LITERAL))
return self._ion.node_by_id[self._ci.ion_node_id]
def ion_name(self):
"""
Return the ION name
:return:
"""
return self._ion.name
def ion_policy(self):
"""
Return the ION policy
:return:
"""
return self._ion.policy
def ion_nodes(self):
"""
Return a list of ION components
:return:
"""
return copy.deepcopy(self._ion.nodes)
def get_ion_node(self, name):
"""
Return a component object given its name
:param name:
:return: Component object matching name (if found), None if not found
"""
self._debug("Getting {} component [{}]".format(Constants.ION_LITERAL, name))
self._debug("{} comps by name: {}".format(Constants.ION_LITERAL, self._ion.node_by_name))
if name in self._ion.node_by_name:
return self._ion.node_by_name[name]
return None
def get_ion_node_by_pipeline_instance(self, pipe_instance_id):
"""
Return a list of ion_nodes for the given pipeline instance.
:param pipe_instance_id: id of the pipeline
:return: ion_nodes or None if not found
"""
if pipe_instance_id in self._ion.node_by_pipe_inst_id:
return self._ion.node_by_pipe_inst_id[pipe_instance_id]
return None
def get_ion_node_agents(self, node):
"""
Return a list of agents for the given ion_node
:param node: a node within the ION
:return: List of agents for the given ion_node
:raises MLOpsException for invalid arguments
"""
# Getting by component name
if isinstance(node, six.string_types):
if node not in self._ion.node_by_name:
raise MLOpsException("Node {} is not part of current {}".format(
node, Constants.ION_LITERAL))
node_obj = self._ion.node_by_name[node]
if node_obj.ee_id not in self._ees_dict:
raise MLOpsException("Component {} had ee_id {} which is not part of valid ees".format(
node_obj.name, node_obj.ee_id))
ee_obj = self._ees_dict[node_obj.ee_id]
# Note: calling deepcopy in order for the user to get a copy of agents objects and not point to internal
# data
agent_list = copy.deepcopy(ee_obj.agents)
return agent_list
else:
raise MLOpsException("component argument should be component name (string)")
def get_agent_by_id(self, agent_id):
"""
Return agent object by ID, assuming the agent is part of the current ION
:param agent_id: Agent Id to search for
:type agent_id: str
:return: agent_id or None if not found
"""
if agent_id in self._agents_dict:
return self._agents_dict[agent_id]
return None
def done(self):
if self._rest_helper is not None:
self._rest_helper.done()