impove the implementation of the wating function. Now we replace rolling as a…

impove the implementation of the wating function. Now we replace rolling as a queue to check if the jobs are finished.

source/train/XLeader.cpp

@@ -118,11 +118,57 @@ void XLeader::SetServerModel(XConfig * config, XModel * model)
 void XLeader::InitForRun()
 {
     serverModel.InitForRun();
-    
+
     for (int i = 0; i < jworkers.count; i++) {
-        XModel * model = ((XWorkerJob*)jworkers[i])->GetModel();
+        XModel* model = ((XWorkerJob*)jworkers[i])->GetModel();
         model->InitForRun();
     }
+
+    XList workers;
+    workers.AddList(&jworkers);
+    workers.AddList(&cworkers);
+    workers.AddList(&uworkers);
+    workers.AddList(&bworkers);
+
+    for (int i = 0; i < workers.count; i++) {
+        XWorker* worker = (XWorker*)workers[i];
+        CheckNTErrors(worker->IsEmpty(), "Something is wrong with the finishedQueue!");
+    }
+}
+
+/*
+wait for finished states (i.e., all workers finish their jobs)
+>> activeJobWorkers - indicates whether each job worker is active
+*/
+void XLeader::WaitForFinishing(const int* activeJobWorkers)
+{
+    int activeCount = 0;
+    for (int i = 0; i < jworkers.count; i++) {
+        if (activeJobWorkers[i] > 0) {
+            XWorker* worker = (XWorker*)jworkers[i];
+            worker->DequeueFinishedJob();
+            activeCount++;
+        }
+    }
+
+    if (activeCount > 0) {
+        for (int i = 0; i < cworkers.count; i++) {
+            XWorker* worker = (XWorker*)cworkers[i];
+            worker->DequeueFinishedJob();
+        }
+
+        for (int i = 0; i < uworkers.count; i++) {
+            XWorker* worker = (XWorker*)uworkers[i];
+            for (int j = 0; j < serverModel.paramNum; j++)
+                worker->DequeueFinishedJob();
+        }
+
+        for (int i = 0; i < bworkers.count; i++) {
+            XWorker* worker = (XWorker*)bworkers[i];
+            for (int j = 0; j < serverModel.paramNum; j++)
+                worker->DequeueFinishedJob();
+        }
+    }
 }
 
 /* get loss */
@@ -278,8 +324,6 @@ bool XLeader::Run(XConfig * config, DataDistributeBase * dataDistributor,
     
     InitForRun();
 
-    //LockWorkers();
-
     for (int i = 0; i < jworkers.count; i++)
         active[i] = 0;
 
@@ -308,6 +352,9 @@ bool XLeader::Run(XConfig * config, DataDistributeBase * dataDistributor,
             /* job in queue 1: make a record of the run */
             worker->AddJobRecord(&serverRecord);
 
+            /* job in queue 1: mark finished */
+            worker->AddJobEnqueueFinished();
+
             active[i] = 1;
             activeJobCount++;
         }
@@ -339,7 +386,8 @@ bool XLeader::Run(XConfig * config, DataDistributeBase * dataDistributor,
 
         collecter->AddJobUpdateAll(&members, &membersAll, &serverModel, 
                                    optimizer, updater, broadcaster);
-        collecter->AddJobCollectOther(&memberRecords, &serverRecord);
+        //collecter->AddJobCollectOther(&memberRecords, &serverRecord);
+        collecter->AddJobEnqueueFinished();
         
         /* jobs in queue 2: collect the (gradient) data and other stuff. This
            is a reduce process. */
@@ -354,9 +402,11 @@ bool XLeader::Run(XConfig * config, DataDistributeBase * dataDistributor,
            not involved in this run. */
         //broadcaster->AddJobBroadcast(&serverModel, &membersAll);
 
-        WaitForFinishing();
+        //WaitForFinishing();
     }
 
+    WaitForFinishing(active);
+
     for (int i = 0; i < jworkers.count; i++) {
         XWorkerJob * worker = (XWorkerJob*)jworkers[i];
         worker->Clear();

source/train/XLeader.h

@@ -112,11 +112,8 @@ public:
     /* initialize the models for running them */
     void InitForRun();
 
-    /* mark the workers as LOCKED */
-    void LockWorkers();
-
-    /* wait for unlocked workers (i.e., all workers finish their jobs) */
-    void WaitForUnlockedWorkers(const int * activeJobWorkers);
+    /* wait for finished states (i.e., all workers finish their jobs) */
+    void WaitForFinishing(const int * activeJobWorkers);
 
     /* get loss */
     float GetLoss();

source/train/XWorker.cpp

@@ -110,4 +110,55 @@ bool XWorker::IsEmpty()
     return queue.IsEmpty();
 }
 
+/* enqueue a counting job of a finished job */
+void XWorker::EnqueueFinishedJob()
+{
+    finishedQueue.Enqueue(NULL);
+}
+
+/* dequeue a counting job of a finished job */
+void XWorker::DequeueFinishedJob()
+{
+    finishedQueue.Dequeue();
+}
+
+/* wrapper of EnqueueFinished() */
+void XWorker::EnqueueFinished(XList* args)
+{
+    XWorker* worker = (XWorker*)args->GetItem(0);
+    worker->EnqueueFinishedJob();
+}
+
+/* wrapper of DequeueFinished() */
+void XWorker::DequeueFinished(XList* args)
+{
+    XWorker* worker = (XWorker*)args->GetItem(0);
+    worker->DequeueFinishedJob();
+}
+
+/* add a job of enqueuing a counting a finished job */
+void XWorker::AddJobEnqueueFinished()
+{
+    XList args;
+    args.Add(this);
+
+    if (isInstantRun)
+        XWorker::EnqueueFinished(&args);
+    else
+        queue.EnqueueJob((void*)(char*)XWorker::EnqueueFinished, &args);
+}
+
+/* add a job of dequeuing a counting a finished job */
+void XWorker::AddJobDequeueFinished()
+{
+    XList args;
+    args.Add(this);
+
+    if (isInstantRun)
+        XWorker::DequeueFinished(&args);
+    else
+        queue.EnqueueJob((void*)(char*)XWorker::DequeueFinished, &args);
+
+}
+
 } /* end of the nts (NiuTrans.Tensor) namespace */

source/train/XWorker.h

@@ -55,7 +55,7 @@ protected:
     /* id of the worker */
     int id;
 
-    /* the queue */
+    /* the queue of jobs */
     XQueue queue;
 
     /* state of the worker */
@@ -63,6 +63,9 @@ protected:
 
     /* fire the flag of instant run */
     bool isInstantRun;
+
+    /* the queue of counting finished jobs */
+    XQueue finishedQueue;
     
 public:
     /* constructor */
@@ -100,6 +103,26 @@ public:
 
     /* whether the job queue is empty? */
     bool IsEmpty();
+
+    /* enqueue a counting job of a finished job */
+    void EnqueueFinishedJob();
+
+    /* dequeue a counting job of a finished job */
+    void DequeueFinishedJob();
+
+    /* wrapper of EnqueueFinished() */
+    static
+    void EnqueueFinished(XList* args);
+
+    /* wrapper of DequeueFinished() */
+    static
+    void DequeueFinished(XList* args);
+
+    /* add a job of enqueuing a counting a finished job */
+    void AddJobEnqueueFinished();
+
+    /* add a job of dequeuing a counting a finished job */
+    void AddJobDequeueFinished();
 };
 
 }

source/train/XWorkerBroadcast.cpp

@@ -223,49 +223,4 @@ bool XWorkerBroadcast::AddJobBroadcast(XModel * source, XList * targetList)
     return true;
 }
 
-
-/* 
-mark the state of the parameter to FINISHED 
->> source - the model that we are updating
->> pid - the parameter index
-*/
-void XWorkerBroadcast::FinishUpdateSingle(XModel * source, int pid)
-{
-    source->params[pid].trainFlag = PARAM_STATE_UPDATED;
-    MUTEX_UNLOCK(source->params[pid].trainLock);
-}
-
-/* wrapper of FinishUpdateSingle */
-void XWorkerBroadcast::FinishSingle(XList * args)
-{
-    XWorkerBroadcast * broadcaster = (XWorkerBroadcast*)args->GetItem(0);
-    XModel * source = (XModel*)args->GetItem(1);
-    int pid = args->GetInt(2);
-
-    broadcaster->FinishUpdateSingle(source, pid);
-}
-
-/* 
-add a new job of finishing the update
->> source - the model that we are updating
->> pid - the parameter index
-*/
-bool XWorkerBroadcast::AddJobFinish(XModel * source, int pid)
-{
-    CheckNTErrors(source != NULL, "no input source tensor!");
-    CheckNTErrors(pid >= 0 && pid < source->paramNum, "illegal parameter index!");
-
-    XList args;
-    args.Add(this);
-    args.Add(source);
-    args.AddInt(pid);
-
-    if (isInstantRun)
-        XWorkerBroadcast::FinishSingle(&args);
-    else
-        queue.EnqueueJob((void*)(char*)XWorkerBroadcast::FinishSingle, &args);
-
-    return true;
-}
-
 }

source/train/XWorkerBroadcast.h

@@ -82,16 +82,6 @@ public:
 
     /* add a new job of broadcasting data (for a model) */
     bool AddJobBroadcast(XModel * source, XList * targetList);
-
-    /* mark the state of the parameter to FINISHED */
-    void FinishUpdateSingle(XModel * source, int pid);
-
-    /* wrapper of FinishUpdateSingle */
-    static
-    void FinishSingle(XList * args);
-
-    /* add a new job of finishing the update */
-    bool AddJobFinish(XModel * source, int pid);
 };
 
 }

source/train/XWorkerCollect.cpp

@@ -126,6 +126,8 @@ void XWorkerCollect::UpdateDataAll(XList * memberActive, XList * memberAll, XMod
                             paramServer.flag = PARAM_STATE_COLLECTED;
                             if (updater != NULL) {
                                 updater->AddJobUpdateSingle(server, memberAll, j, optimizer, broadcaster);
+                                updater->AddJobEnqueueFinished();
+
                             }
                         }
                         else if (finishedCount[j] > memberActive->count) {
@@ -135,63 +137,6 @@ void XWorkerCollect::UpdateDataAll(XList * memberActive, XList * memberAll, XMod
                 }
             }
         }
-        else if (collectMode == DATA_COLLECT_REDUCESUM) {
-            for (int j = 0; j < server->paramNum; j++) {
-
-                bool ready = true;
-                XParamKeeper &paramServer = server->params[j];
-
-                /* tp[j]->isGradFinished is true only if the model finishes the computation
-                (in another process) */
-                if (paramServer.flag != PARAM_STATE_NOT_READY || !paramServer.param->isGradFinished)
-                    continue;
-
-                /* check if all the models (or part of them) are ready */
-                for (int i = 0; i < memberActive->count; i++) {
-                    XModel * source = (XModel*)memberActive->GetItem(i);
-                    XParamKeeper &paramSource = source->params[j];
-
-                    /* sp[j]->isGradFinished is true only if the model finishes the computation
-                    (in another process) */
-                    if (paramSource.flag == PARAM_STATE_COLLECTED ||
-                        paramSource.flag == PARAM_STATE_UPDATED ||
-                        !paramSource.param->isGradFinished)
-                    {
-                        ready = false;
-                        break;
-                    }
-                    else if (paramSource.flag == PARAM_STATE_NOT_READY) {
-                        paramSource.flag = PARAM_STATE_READY;
-                    }
-                }
-
-                if (ready) {
-                    XList tensorList(memberActive->count);
-
-                    for (int i = 0; i < memberActive->count; i++) {
-                        XModel * source = (XModel*)memberActive->GetItem(i);
-                        tensorList.Add(source->params[j].param->grad);
-                    }
-
-                    /* data transmit */
-                    CollectReduceSum(&tensorList, server->params[j].param->grad);
-
-                    /* reset the flags */
-                    for (int i = 0; i < memberActive->count; i++) {
-                        XModel * source = (XModel*)memberActive->GetItem(i);
-                        source->params[j].flag = PARAM_STATE_COLLECTED;
-                    }
-
-                    server->params[j].flag = PARAM_STATE_COLLECTED;
-                    finished += memberActive->count;
-
-                    /* we call model update (in another thread) and then
-                       broadcast the new parameters to member models
-                       (in another thread) */
-                    updater->AddJobUpdateSingle(server, memberAll, j, optimizer, broadcaster);
-                }
-            }
-        }
         else {
             ShowNTErrors("Unsupported data collection mode!");
         }

source/train/XWorkerUpdate.cpp

@@ -79,6 +79,7 @@ void XWorkerUpdate::UpdateParameter(XModel * server, XList * members, int pid,
 
     /* broadcast the new parameter to other models (in anotehr worker/thread) */
     broadcaster->AddJobBroadcastSingle(server, members, pid);
+    broadcaster->AddJobEnqueueFinished();
 }
 
 /*