def test_register_models(self): try: os.remove(DB.database) except FileNotFoundError: pass MLModel.create_table() for x in range(1, 101): job_id = str(time.time()) model = MLModel(f_role='host', f_party_id='100', f_job_id=job_id, f_model_id=f'foobar#{x}', f_model_version=job_id, f_initiator_role='host', f_work_mode=0) model.save(force_insert=True) self.assertEqual( db_services.models_group_by_party_model_id_and_model_version(). count(), 100) with patch.object(self.service_db, '_insert') as insert: self.service_db.register_models() self.assertEqual(insert.call_count, 100) with patch.object(self.service_db, '_delete') as delete: self.service_db.unregister_models() self.assertEqual(delete.call_count, 100) os.remove(DB.database)
def save_model_info(model_info): model = MLModel() model.f_create_time = current_timestamp() for k, v in model_info.items(): attr_name = 'f_%s' % k if hasattr(MLModel, attr_name): setattr(model, attr_name, v) elif hasattr(MLModel, k): setattr(model, k, v) try: rows = model.save(force_insert=True) if rows != 1: raise Exception("Create {} failed".format(MLModel)) if RuntimeConfig.zk_client is not None: ServiceUtils.register( RuntimeConfig.zk_client, gen_party_model_id(role=model.f_role, party_id=model.f_party_id, model_id=model.f_model_id), model.f_model_version) return model except peewee.IntegrityError as e: if e.args[0] == 1062: sql_logger(job_id=model_info.get("job_id", "fate_flow")).warning(e) else: raise Exception("Create {} failed:\n{}".format(MLModel, e)) except Exception as e: raise Exception("Create {} failed:\n{}".format(MLModel, e))
def save_model_info(model_info): model = MLModel() model.f_create_time = current_timestamp() for k, v in model_info.items(): attr_name = 'f_%s' % k if hasattr(MLModel, attr_name): setattr(model, attr_name, v) elif hasattr(MLModel, k): setattr(model, k, v) try: rows = model.save(force_insert=True) if rows != 1: raise Exception("Save to database failed") except peewee.IntegrityError as e: if e.args[0] != 1062: raise Exception("Create {} failed:\n{}".format(MLModel, e)) sql_logger(job_id=model_info.get("job_id", "fate_flow")).warning(e) return except Exception as e: raise Exception("Create {} failed:\n{}".format(MLModel, e)) RuntimeConfig.SERVICE_DB.register_model( gen_party_model_id(role=model.f_role, party_id=model.f_party_id, model_id=model.f_model_id), model.f_model_version) return model
def save_model_info(model_info): model = MLModel() model.f_create_time = current_timestamp() for k, v in model_info.items(): attr_name = 'f_%s' % k if hasattr(MLModel, attr_name): setattr(model, attr_name, v) elif hasattr(MLModel, k): setattr(model, k, v) try: rows = model.save(force_insert=True) if rows != 1: raise Exception("Create {} failed".format(MLModel)) return model except peewee.IntegrityError as e: if e.args[0] == 1062: sql_logger(job_id=model_info.get("job_id", "fate_flow")).warning(e) else: raise Exception("Create {} failed:\n{}".format(MLModel, e)) except Exception as e: raise Exception("Create {} failed:\n{}".format(MLModel, e))