The pserver now (optionally) prunes the provDB

To enable the above, the pserver now requires the full algo_params JSON file to be passed in vs just the algorithm name.
Reordered the arguments to PSshardProvenanceDBclient::connectShard to put shard index first
Fixed bug in AlgoParams JSON parse where the checking of present arguments was being performed incorrectly for algorithms other than the HBOS default

Author: giltirn

app/pserver.cpp

@@ -10,13 +10,15 @@
 
 #ifdef ENABLE_PROVDB
 #include <chimbuko/core/pserver/PSglobalProvenanceDBclient.hpp>
+#include <chimbuko/core/pserver/PSshardProvenanceDBclient.hpp>
 #endif
 
 #include <chimbuko/core/param/sstd_param.hpp>
 #include <chimbuko/core/util/commandLineParser.hpp>
 #include <chimbuko/core/util/error.hpp>
 #include <fstream>
-#include "chimbuko/core/verbose.hpp"
+#include <chimbuko/core/verbose.hpp>
+#include <chimbuko/core/ad/ADOutlier.hpp>
 
 #include <chimbuko/modules/factory.hpp>
 
@@ -40,11 +42,12 @@ struct pserverArgs{
   int stat_send_freq;
 
   std::string stat_outputdir;
-  std::string ad;
 
   int model_update_freq; //frequency in ms at which the global model is updated
   bool model_force_update; //force the global model to be updated every time a worker thread updates its model
 
+  ADOutlier::AlgoParams algo_params; //The AD algorithm hyperparameters and type/name
+
 #ifdef _USE_ZMQNET
   int max_pollcyc_msg;
   int zmq_io_thr;
@@ -53,25 +56,27 @@ struct pserverArgs{
 
 #ifdef ENABLE_PROVDB
   std::string provdb_addr_dir;
+  int nprovdb_shards; /**< Number of database shards*/
+  int nprovdb_instances; /**< Number of instances of the provenance database server*/
   std::string provdb_mercury_auth_key; //An authorization key for initializing Mercury (optional, default "")
+  bool provdb_post_prune; //perform post-pruning on the provenance database
 #endif
 
   std::string prov_outputpath;
 
-  pserverArgs(): ad("hbos"), nt(-1), logdir("."), ws_addr(""), load_params_set(false), save_params_set(false), freeze_params(false), stat_send_freq(1000), stat_outputdir(""), port(5559), prov_outputpath(""), model_update_freq(1000), model_force_update(false)
+  pserverArgs(): nt(-1), logdir("."), ws_addr(""), load_params_set(false), save_params_set(false), freeze_params(false), stat_send_freq(1000), stat_outputdir(""), port(5559), prov_outputpath(""), model_update_freq(1000), model_force_update(false)
 #ifdef _USE_ZMQNET
 	       , max_pollcyc_msg(10), zmq_io_thr(1), autoshutdown(true)
 #endif
 #ifdef ENABLE_PROVDB
-	       , provdb_addr_dir(""), provdb_mercury_auth_key("")
+	       , provdb_addr_dir(""), provdb_mercury_auth_key(""), provdb_post_prune(true), nprovdb_shards(1), nprovdb_instances(1)
 #endif
   {}
 
   static commandLineParser &getParser(pserverArgs &instance){
     static bool init = false;
     static commandLineParser p;
     if(!init){
-      addOptionalCommandLineArg(p, instance, ad, "Set AD algorithm to use.");
       addOptionalCommandLineArg(p, instance, nt, "Set the number of RPC handler threads (max-2 by default)");
       addOptionalCommandLineArg(p, instance, logdir, "Set the output log directory (default: job directory)");
       addOptionalCommandLineArg(p, instance, port, "Set the pserver port (default: 5559)");
@@ -89,11 +94,15 @@ struct pserverArgs{
 #ifdef ENABLE_PROVDB
       addOptionalCommandLineArg(p, instance, provdb_addr_dir, "The directory containing the address file written out by the provDB server. An empty string will disable the connection to the global DB.  (default empty, disabled)");
       addOptionalCommandLineArg(p, instance, provdb_mercury_auth_key, "Set the Mercury authorization key for connection to the provDB (default \"\")");
+      addOptionalCommandLineArg(p, instance, provdb_post_prune, "If enabled the pserver will automatically \"prune\" the provenance database at the end of the run (default: true)");
+      addOptionalCommandLineArgWithDefault(p, instance, nprovdb_shards, 1, "Number of provenance database shards. Clients connect to shards round-robin by rank (default 1)");
+      addOptionalCommandLineArgWithDefault(p, instance, nprovdb_instances, 1, "Number of provenance database instances. Shards are divided uniformly over instances. (default 1)");
 #endif
       addOptionalCommandLineArg(p, instance, prov_outputpath, "Output global provenance data to this directory. Can be used in place of or in conjunction with the provenance database. An empty string \"\" (default) disables this output");
       addOptionalCommandLineArg(p, instance, model_update_freq, "The frequency in ms at which the global AD model is updated (default 1000ms)");
       addOptionalCommandLineArg(p, instance, model_force_update, "Force the global AD model to be updated every time a worker thread updates its model (default false)");
 
+      p.addOptionalArg(new ADOutlier::AlgoParams::cmdlineParser(instance.algo_params, "-algo_params_file", "Set the filename containing the algorithm name and hyperparameters (ensure consistent with OAD)."));
       init = true;
     }
     return p;
@@ -132,7 +141,7 @@ int main (int argc, char ** argv){
     enableVerboseLogging() = true;
   }
 
-  PSparamManager param(args.nt, args.ad); //the AD model; independent models for each worker thread that are aggregated periodically to a global model
+  PSparamManager param(args.nt, args.algo_params.algorithm); //the AD model; independent models for each worker thread that are aggregated periodically to a global model
   param.enableForceUpdate(args.model_force_update); //decide whether the model is forced to be updated every time a worker updates its model
 
   std::unique_ptr<ProvDBmoduleSetupCore> pdb_setup = modules::factoryInstantiateProvDBmoduleSetup(module);
@@ -272,6 +281,22 @@ int main (int argc, char ** argv){
     ps_module_data_man->writeModel(args.save_params, param);
   }
 
+
+  //Post-prune the provenance database
+  if(args.provdb_post_prune && args.provdb_addr_dir.size()){
+    progressStream << "PServer: Pruning the provenance database" << std::endl;
+    std::unique_ptr<ProvDBpruneCore> pruner = modules::factoryInstantiateProvDBprune(module, args.algo_params, param.getGlobalParamsCopy()->serialize());
+
+    for(int s=0;s<args.nprovdb_shards;s++){
+      progressStream << "PServer: Pruning shard " << s+1 << " of " << args.nprovdb_shards << std::endl;
+      PSshardProvenanceDBclient shard_client(pdb_setup->getMainDBcollections());
+      shard_client.connectShard(s, args.provdb_addr_dir, args.nprovdb_shards, args.nprovdb_instances);
+      pruner->prune(shard_client.getDatabase());
+    }
+    progressStream << "PServer: Updating global function stats" << std::endl;
+    pruner->finalize(provdb_client.getDatabase());
+  }
+
   progressStream << "Pserver: finished" << std::endl;
 
   return 0;

include/chimbuko/core/provdb/setup.hpp

@@ -74,6 +74,11 @@ namespace chimbuko{
      */
     inline int getNshards() const{ return m_nshards; }
 
+    /**
+     * @brief Return the total number of instances
+     */
+    inline int getNinstance() const{ return m_ninstances; }
+
     /**
      * @brief Get the number of shards serviced by a given instance
      */

include/chimbuko/core/pserver/PSshardProvenanceDBclient.hpp

@@ -21,7 +21,7 @@ namespace chimbuko{
      * @brief Connect the client the appropriate provenance database server instance / shard using the default setup
      * @param addr_file_dir The directory containing the address files created by the provDB
      */
-    void connectShard(const std::string &addr_file_dir, int shard, int nshards, int ninstances);
+    void connectShard(int shard, const std::string &addr_file_dir, const int nshards, const int ninstances);
 
     /**
      * @brief No handshake is needed

scripts/launch/run_services.sh

@@ -204,7 +204,7 @@ if (( ${use_provdb} == 1 )); then
     done
 
     extra_args+=" -provdb_addr_dir ${provdb_addr_dir} -nprovdb_instances ${provdb_ninstances} -nprovdb_shards ${provdb_nshards}"
-    ps_extra_args+=" -provdb_addr_dir ${provdb_addr_dir}"
+    ps_extra_args+=" -provdb_addr_dir ${provdb_addr_dir} -nprovdb_instances ${provdb_ninstances} -nprovdb_shards ${provdb_nshards}"
     echo "Chimbuko Services: Enabling provenance database with arg: ${extra_args}"
     cd -
 
@@ -370,7 +370,7 @@ if (( ${use_pserver} == 1 )); then
 
     pserver_alg=${ad_alg} #Pserver AD algorithm choice must match that used for the driver
     pserver_addr="tcp://${ip}:${pserver_port}"  #address for parameter server in format "tcp://IP:PORT"
-    cmd="pserver ${module} -ad ${pserver_alg} -nt ${pserver_nt} -logdir ${log_dir} -port ${pserver_port} -save_params ${ps_dir}/global_model.json ${ps_extra_args} 2>&1 | tee ${log_dir}/pserver.log  &"
+    cmd="pserver ${module} -algo_params_file ${algo_params_file} -nt ${pserver_nt} -logdir ${log_dir} -port ${pserver_port} -save_params ${ps_dir}/global_model.json ${ps_extra_args} 2>&1 | tee ${log_dir}/pserver.log  &"
     if [[ ! -z "${pserver_numa_bind:-}" ]]; then
 	echo "Chimbuko Services binding pserver to NUMA domain ${pserver_numa_bind}"
 	cmd="numactl -N ${pserver_numa_bind} ${cmd}"

src/core/ad/ADOutlier.cpp

@@ -23,6 +23,8 @@ void ADOutlier::AlgoParams::setJson(const nlohmann::json &in){
 #define JSON_GET(to) if(in.contains(#to)) to = in[#to].template get<decltype(to)>()
   //Check for required
   JSON_CHECK(algorithm);
+  JSON_GET(algorithm);
+
   if(algorithm == "sstd"){
     JSON_CHECK(sstd_sigma);
   }else if(algorithm == "hbos"){
@@ -33,7 +35,6 @@ void ADOutlier::AlgoParams::setJson(const nlohmann::json &in){
     JSON_CHECK(hbos_thres);
   }
   //Get all available
-  JSON_GET(algorithm);
   JSON_GET(sstd_sigma);
   JSON_GET(glob_thres);
   JSON_GET(hbos_max_bins);
@@ -62,9 +63,7 @@ nlohmann::json ADOutlier::AlgoParams::getJson() const{
 }
 
 int ADOutlier::AlgoParams::cmdlineParser::parse(const std::string &arg, const char** vals, const int vals_size){
-  std::cout << "TEST ARG " << arg << "==" << m_arg << std::endl;
   if(arg == m_arg){
-    std::cout << "TEST ARG FOUND ARG " << arg << std::endl;
     if(vals_size < 1) return -1;
 
     try{

src/core/pserver/PSshardProvenanceDBclient.cpp

@@ -8,7 +8,7 @@ using namespace chimbuko;
 
 PSshardProvenanceDBclient::~PSshardProvenanceDBclient(){ disconnect(); } //call disconnect in derived class to ensure derived class is still alive when disconnect is called
 
-void PSshardProvenanceDBclient::connectShard(const std::string &addr_file_dir, int shard, int nshards, int ninstances){
+void PSshardProvenanceDBclient::connectShard(int shard, const std::string &addr_file_dir, const int nshards, const int ninstances){
   ProvDBsetup setup(nshards, ninstances);
 
   int instance = setup.getShardInstance(shard);

src/modules/performance_analysis/provdb/ProvDBprune.cpp

@@ -1,6 +1,7 @@
 #include<chimbuko/modules/performance_analysis/provdb/ProvDBprune.hpp>
 #include<chimbuko/modules/performance_analysis/provdb/ProvDBpruneInterface.hpp>
 #include<chimbuko/core/util/error.hpp>
+#include<chimbuko/core/verbose.hpp>
 #include<limits>
 
 using namespace chimbuko;
@@ -10,11 +11,13 @@ void ProvDBprune::pruneImplementation(ADOutlier &ad, sonata::Database &db){
   //Prune the outliers and update scores / model on remaining. Also gather new anomaly statistics
   {
     ProvDBpruneInterface po(ad, db, ADDataInterface::EventType::Outlier, &m_anom_metrics);
+    progressStream << "Pruning normal events from 'anomalies' database with " << po.nDataSets() << " function indices" << std::endl;
     ad.run(po);
   }
   //Prune normal execs and update scores / model on remaining
   {
     ProvDBpruneInterface pn(ad, db, ADDataInterface::EventType::Normal);
+    progressStream << "Pruning outlier events from 'normal_execs' database with " << pn.nDataSets() << " function indices" << std::endl;
     ad.run(pn);
   }
 }

src/modules/performance_analysis/provdb/ProvDBpruneInterface.cpp

@@ -26,10 +26,7 @@ while( ($rec = db_fetch(')" + coll + R"(')) != NULL ){
 
   nlohmann::json rj = nlohmann::json::parse(r->second);
   if(!rj.is_array()) fatal_error("Expected an array type");
-  
-  for(int i=0;i<rj.size();i++)
-    std::cout << i << " " << rj[i][0] << " " << rj[i][1] << " " << rj[i][2] << std::endl;
-  
+   
   for(auto const &e: rj){
     uint64_t id = e[0].template get<uint64_t>();
     unsigned long fid = e[1].template get<unsigned long>();
@@ -54,6 +51,7 @@ void ProvDBpruneInterface::recordDataSetLabelsInternal(const std::vector<Elem> &
   std::vector<uint64_t> to_update;
   std::vector<double> update_scores;
 
+  int fid = this->getDataSetModelIndex(dset_index);
   ADDataInterface::EventType type_to_remove = m_prune_type == ADDataInterface::EventType::Outlier ? ADDataInterface::EventType::Normal : ADDataInterface::EventType::Outlier;
 
   for(auto const &e: data){ //we used the record id as the index, which is unique
@@ -65,12 +63,12 @@ void ProvDBpruneInterface::recordDataSetLabelsInternal(const std::vector<Elem> &
       update_scores.push_back(e.score);
     }
   }
-  progressStream << "Pruning " << to_prune.size() << " of " << data.size() << " records in shard" << std::endl;
-
+  progressStream << "Pruning " << to_prune.size() << " of " << data.size() << " records in shard for function index " << fid << std::endl;
+  
   m_collection->erase_multi(to_prune.data(), to_prune.size(), true); //last entry tells database to commit change
 
   //Grab records to update in batches
-  progressStream << "Updating scores and models for " << to_update.size() << " records" << std::endl;
+  progressStream << "Updating scores and models for " << to_update.size() << " records with function index " << fid << std::endl;
   std::unordered_map<uint64_t, nlohmann::json> json_param_cache;
   batchAmendRecords(*m_collection, to_update, [&](nlohmann::json &rec, size_t i){  
       //Update score

test/run_ad.sh

@@ -21,7 +21,16 @@ mkdir -p temp perf
 
 export CHIMBUKO_DISABLE_CUDA_JIT_WORKAROUND=1
 
-mpirun --allow-run-as-root --oversubscribe -n 1 ${appdir}/pserver performance_analysis -logdir "./perf/" -ad "sstd" -model_force_update 1 &
+cat <<EOF > /tmp/algo_params.json
+{
+  "algorithm" : "sstd",
+  "sstd_sigma" : 6.0
+}
+EOF
+
+cat /tmp/algo_params.json
+echo "Starting pserver"
+mpirun --allow-run-as-root --oversubscribe -n 1 ${appdir}/pserver performance_analysis -logdir "./perf/" -algo_params_file "/tmp/algo_params.json" -model_force_update 1 &
 ps_wid=$!
 
 sleep 5

test/unit_tests/modules/performance_analysis/provdb/ProvDBprune.cpp

@@ -22,7 +22,7 @@ TEST(TestProvDBprune, works){
   std::vector<std::unique_ptr<PSshardProvenanceDBclient> > shard_clients(nshards);
   for(int i=0;i<nshards;i++){
     shard_clients[i].reset(new PSshardProvenanceDBclient(setup.getMainDBcollections()));
-    shard_clients[i]->connectShard("/tmp",i,nshards,1);
+    shard_clients[i]->connectShard(i,"/tmp",nshards,1);
   }
 
   PSglobalProvenanceDBclient glob_client(setup.getGlobalDBcollections());