trainer.h

/************************************************************
*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements.  See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership.  The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License.  You may obtain a copy of the License at
* 
*   http://www.apache.org/licenses/LICENSE-2.0
* 
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied.  See the License for the
* specific language governing permissions and limitations
* under the License.
*
*************************************************************/

#ifndef INCLUDE_TRAINER_TRAINER_H_
#define INCLUDE_TRAINER_TRAINER_H_
#include <unordered_map>
#include <queue>
#include "proto/job.pb.h"
#include "proto/singa.pb.h"
#include "utils/param.h"
#include "utils/singleton.h"
#include "utils/factory.h"
#include "neuralnet/neuralnet.h"
#include "trainer/worker.h"
#include "trainer/server.h"
#include "communication/socket.h"

namespace singa {
class Trainer{
 public:
  ~Trainer();
  void Start(bool resume, const SingaProto& singaConf, JobProto* jobConf);

 protected:
  void Resume(JobProto* jobConf);
  vector<Server*> CreateServers(int nthread, const JobProto& jobConf);
  vector<Worker*> CreateWorkers(int nthread, const JobProto& jobConf);

  void SetupWorkerServer(
    const JobProto& jobConf,
    const vector<Worker*>& workers,
    const vector<Server*>& servers);

  void Run(const vector<Worker*>& workers, const vector<Server*>& servers);
  void DisplayMetric(Msg** msg);
  Dealer* CreateInterProcsDealer(int dst_procs);
  void HandleLocalMsg(std::queue<Msg*>* msg_queue, Msg** msg);

    const vector<Msg*> HandleGet(ParamEntry* entry, Msg** msg);
    void HandleGetResponse(ParamEntry* entry, Msg** msg);

    const vector<Msg*> HandleUpdate(ParamEntry* entry, Msg** msg);
  void HandleUpdateResponse(ParamEntry* entry, Msg** msg);

    const vector<Msg*> HandlePut(ParamEntry* entry, Msg** msg);

  void GenMsgs(int type, int version, ParamEntry* entry,
    Msg* msg, vector<Msg*> *ret);
  inline int Hash(int grp_id, int param_id) {
    return grp_id * 997 + param_id;
  }

 protected:
  int procs_id_;
  Router *router_;
  std::unordered_map<int, ParamEntry*> worker_shard_;
  vector<int> slice2server_;
};
} /* singa */
#endif // INCLUDE_TRAINER_TRAINER_H_