<commit_before>package bwmf

import (
	"fmt"
	"log"
	"math/rand"
	"os"
	"strconv"
	"time"

	"github.com/golang/protobuf/proto"
	"github.com/taskgraph/taskgraph"
	pb "github.com/taskgraph/taskgraph/example/bwmf/proto"
	"github.com/taskgraph/taskgraph/filesystem"
	"github.com/taskgraph/taskgraph/op"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
)

/*
The block wise matrix factorization task is designed for carry out block wise matrix
factorization for a variety of criteria (loss function) and constraints (non-negativity
for example).

The main idea behind the bwmf is following:
We will have K tasks that handle both row task and column task in alternation. Each task
will read two copies of the data: one row shard and one column shard of A. It either hosts
one shard of D and a full copy of T, or one shard of T and a full copy of D, depending on
the epoch of iteration. "A full copy" consists of computation results from itself and
all "children".
Topology: the topology is different from task to task. Each task will consider itself parent
and all others children.
*/

// bwmfTasks holds two shards of original matrices (row and column), one shard of D,
// and one shard of T. It works differently for odd and even epoch:
// During odd epoch, 1. it fetch all T from other slaves, and finding better value for
// local shard of D; 2. after it is done, it let every one knows. Vice versa for even epoch.
// Task 0 will monitor the progress and responsible for starting the work of new epoch.
type bwmfTask struct {
	framework  taskgraph.Framework
	epoch      uint64
	taskID     uint64
	logger     *log.Logger
	numOfTasks uint64
	curIter    uint64

	rowShard    *pb.MatrixShard
	columnShard *pb.MatrixShard
	dShard      *pb.MatrixShard
	tShard      *pb.MatrixShard

	peerShards  map[uint64]*pb.MatrixShard
	peerUpdated map[uint64]bool

	dims   *dimensions
	config *Config

	fsClient filesystem.Client

	// optimization toolkits
	dLoss        *KLDivLoss
	tLoss        *KLDivLoss
	dParam       op.Parameter
	tParam       op.Parameter
	optimizer    *op.ProjectedGradient
	stopCriteria op.StopCriteria

	// event hanlding
	epochChange chan *event
	getT        chan *event
	getD        chan *event
	dataReady   chan *event
	updateDone  chan *event
	metaReady   chan *event
	exitChan    chan *event
}

type event struct {
	ctx     context.Context
	epoch   uint64
	request *pb.Request
	retT    chan *pb.Response
	retD    chan *pb.Response
	fromID  uint64
	method  string
	output  proto.Message
}

type dimensions struct {
	m, n, k int
	M, N    int
}

func (t *bwmfTask) initData() {
	var rsErr, csErr error

	rowShardPath := fmt.Sprintf("%s-%06d", t.config.IOConf.IDPath, t.taskID)
	columnShardPath := fmt.Sprintf("%s-%06d", t.config.IOConf.ITPath, t.taskID)
	t.rowShard, rsErr = LoadMatrixShard(t.fsClient, rowShardPath)
	if rsErr != nil {
		t.logger.Panicf("Failed load rowShard from %s with error %s", rsErr)
	}
	t.columnShard, csErr = LoadMatrixShard(t.fsClient, columnShardPath)
	if csErr != nil {
		t.logger.Panicf("Failed load columnShard from %s with error %s", csErr)
	}

	t.dims = &dimensions{
		m: len(t.rowShard.Row),
		n: len(t.columnShard.Row),
		k: t.config.OptConf.DimLatent,
		M: -1,
		N: -1,
	}

	if t.config.IOConf.InitDPath != "" {
		t.dShard, rsErr = LoadMatrixShard(t.fsClient, t.config.IOConf.InitDPath+"."+strconv.Itoa(int(t.taskID)))
		if rsErr != nil {
			t.logger.Panicf("Failed initialize dShard. %s", rsErr)
		}
		if len(t.dShard.Row) != t.dims.m {
			t.logger.Panic("Dimension mismatch for pre-initialized dShard.")
		}
	}

	if t.config.IOConf.InitTPath != "" {
		t.tShard, rsErr = LoadMatrixShard(t.fsClient, t.config.IOConf.InitTPath+"."+strconv.Itoa(int(t.taskID)))
		if rsErr != nil {
			t.logger.Panicf("Failed initialize tShard. %s", rsErr)
		}
		if len(t.tShard.Row) != t.dims.n {
			t.logger.Panic("Dimension mismatch for pre-initialized tShard.")
		}
	}
}

func (t *bwmfTask) initOptUtil() {
	// init optimization utils
	tParamLen := t.dims.n * t.dims.k
	dParamLen := t.dims.m * t.dims.k

	t.tParam = op.NewVecParameter(tParamLen)
	t.dParam = op.NewVecParameter(dParamLen)

	if t.tShard == nil {
		// XXX: Initialize it random and sparse.
		t.tShard = &pb.MatrixShard{
			Row: make([]*pb.MatrixShard_RowData, t.dims.n),
		}
		m1 := int32((10 + t.dims.n) / 10)
		for r, _ := range t.tShard.Row {
			t.tShard.Row[r] = &pb.MatrixShard_RowData{RowId: int32(r), At: make(map[int32]float32)}
			if rand.Int31()%m1 == 0 {
				for k := 0; k < t.dims.k; k++ {
					t.tShard.Row[r].At[int32(k)] = rand.Float32()
				}
			}
		}
	}
	// assign loaded tshard to tparam
	for r, m := range t.tShard.GetRow() {
		for c, v := range m.At {
			t.tParam.Set(r*t.dims.k+int(c), v)
		}
	}

	if t.dShard == nil {
		// XXX: Initial at 0.0.
		t.dShard = &pb.MatrixShard{
			Row: make([]*pb.MatrixShard_RowData, t.dims.m),
		}

		m2 := int32((10 + t.dims.m) / 10)
		for r, _ := range t.dShard.Row {
			t.dShard.Row[r] = &pb.MatrixShard_RowData{RowId: int32(r), At: make(map[int32]float32)}
			if rand.Int31()%m2 == 0 {
				for k := 0; k < t.dims.k; k++ {
					t.dShard.Row[r].At[int32(k)] = rand.Float32()
				}
			}
		}
	}
	// assign loaded dshard to dparam
	for r, m := range t.dShard.GetRow() {
		for c, v := range m.At {
			t.dParam.Set(r*t.dims.k+int(c), v)
		}
	}

	projLen := 0
	if tParamLen > dParamLen {
		projLen = tParamLen
	} else {
		projLen = dParamLen
	}

	t.optimizer = op.NewProjectedGradient(
		op.NewProjection(
			op.NewAllTheSameParameter(1e20, projLen),<commit_msg>bwmf: Remove sparse initialization mechanism<commit_after>package bwmf

import (
	"fmt"
	"log"
	"math/rand"
	"os"
	"strconv"
	"time"

	"github.com/golang/protobuf/proto"
	"github.com/taskgraph/taskgraph"
	pb "github.com/taskgraph/taskgraph/example/bwmf/proto"
	"github.com/taskgraph/taskgraph/filesystem"
	"github.com/taskgraph/taskgraph/op"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
)

/*
The block wise matrix factorization task is designed for carry out block wise matrix
factorization for a variety of criteria (loss function) and constraints (non-negativity
for example).

The main idea behind the bwmf is following:
We will have K tasks that handle both row task and column task in alternation. Each task
will read two copies of the data: one row shard and one column shard of A. It either hosts
one shard of D and a full copy of T, or one shard of T and a full copy of D, depending on
the epoch of iteration. "A full copy" consists of computation results from itself and
all "children".
Topology: the topology is different from task to task. Each task will consider itself parent
and all others children.
*/

// bwmfTasks holds two shards of original matrices (row and column), one shard of D,
// and one shard of T. It works differently for odd and even epoch:
// During odd epoch, 1. it fetch all T from other slaves, and finding better value for
// local shard of D; 2. after it is done, it let every one knows. Vice versa for even epoch.
// Task 0 will monitor the progress and responsible for starting the work of new epoch.
type bwmfTask struct {
	framework  taskgraph.Framework
	epoch      uint64
	taskID     uint64
	logger     *log.Logger
	numOfTasks uint64
	curIter    uint64

	rowShard    *pb.MatrixShard
	columnShard *pb.MatrixShard
	dShard      *pb.MatrixShard
	tShard      *pb.MatrixShard

	peerShards  map[uint64]*pb.MatrixShard
	peerUpdated map[uint64]bool

	dims   *dimensions
	config *Config

	fsClient filesystem.Client

	// optimization toolkits
	dLoss        *KLDivLoss
	tLoss        *KLDivLoss
	dParam       op.Parameter
	tParam       op.Parameter
	optimizer    *op.ProjectedGradient
	stopCriteria op.StopCriteria

	// event hanlding
	epochChange chan *event
	getT        chan *event
	getD        chan *event
	dataReady   chan *event
	updateDone  chan *event
	metaReady   chan *event
	exitChan    chan *event
}

type event struct {
	ctx     context.Context
	epoch   uint64
	request *pb.Request
	retT    chan *pb.Response
	retD    chan *pb.Response
	fromID  uint64
	method  string
	output  proto.Message
}

type dimensions struct {
	m, n, k int
	M, N    int
}

func (t *bwmfTask) initData() {
	var rsErr, csErr error

	rowShardPath := fmt.Sprintf("%s-%06d", t.config.IOConf.IDPath, t.taskID)
	columnShardPath := fmt.Sprintf("%s-%06d", t.config.IOConf.ITPath, t.taskID)
	t.rowShard, rsErr = LoadMatrixShard(t.fsClient, rowShardPath)
	if rsErr != nil {
		t.logger.Panicf("Failed load rowShard from %s with error %s", rsErr)
	}
	t.columnShard, csErr = LoadMatrixShard(t.fsClient, columnShardPath)
	if csErr != nil {
		t.logger.Panicf("Failed load columnShard from %s with error %s", csErr)
	}

	t.dims = &dimensions{
		m: len(t.rowShard.Row),
		n: len(t.columnShard.Row),
		k: t.config.OptConf.DimLatent,
		M: -1,
		N: -1,
	}

	if t.config.IOConf.InitDPath != "" {
		t.dShard, rsErr = LoadMatrixShard(t.fsClient, t.config.IOConf.InitDPath+"."+strconv.Itoa(int(t.taskID)))
		if rsErr != nil {
			t.logger.Panicf("Failed initialize dShard. %s", rsErr)
		}
		if len(t.dShard.Row) != t.dims.m {
			t.logger.Panic("Dimension mismatch for pre-initialized dShard.")
		}
	}

	if t.config.IOConf.InitTPath != "" {
		t.tShard, rsErr = LoadMatrixShard(t.fsClient, t.config.IOConf.InitTPath+"."+strconv.Itoa(int(t.taskID)))
		if rsErr != nil {
			t.logger.Panicf("Failed initialize tShard. %s", rsErr)
		}
		if len(t.tShard.Row) != t.dims.n {
			t.logger.Panic("Dimension mismatch for pre-initialized tShard.")
		}
	}
}

func (t *bwmfTask) initOptUtil() {
	// init optimization utils
	tParamLen := t.dims.n * t.dims.k
	dParamLen := t.dims.m * t.dims.k

	t.tParam = op.NewVecParameter(tParamLen)
	t.dParam = op.NewVecParameter(dParamLen)

	if t.tShard == nil {
		// XXX: Initialize it random and sparse.
		t.tShard = &pb.MatrixShard{
			Row: make([]*pb.MatrixShard_RowData, t.dims.n),
		}
		for r, _ := range t.tShard.Row {
			t.tShard.Row[r] = &pb.MatrixShard_RowData{RowId: int32(r), At: make(map[int32]float32)}
			for k := 0; k < t.dims.k; k++ {
				t.tShard.Row[r].At[int32(k)] = rand.Float32()
			}
		}
	}
	// assign loaded tshard to tparam
	for r, m := range t.tShard.GetRow() {
		for c, v := range m.At {
			t.tParam.Set(r*t.dims.k+int(c), v)
		}
	}

	if t.dShard == nil {
		// XXX: Initial at 0.0.
		t.dShard = &pb.MatrixShard{
			Row: make([]*pb.MatrixShard_RowData, t.dims.m),
		}

		for r, _ := range t.dShard.Row {
			t.dShard.Row[r] = &pb.MatrixShard_RowData{RowId: int32(r), At: make(map[int32]float32)}
			for k := 0; k < t.dims.k; k++ {
				t.dShard.Row[r].At[int32(k)] = rand.Float32()
			}
		}
	}
	// assign loaded dshard to dparam
	for r, m := range t.dShard.GetRow() {
		for c, v := range m.At {
			t.dParam.Set(r*t.dims.k+int(c), v)
		}
	}

	projLen := 0
	if tParamLen > dParamLen {
		projLen = tParamLen
	} else {
		projLen = dParamLen
	}

	t.optimizer = op.NewProjectedGradient(
		op.NewProjection(
			op.NewAllTheSameParameter(1e20, projLen),
{'fraction_non_alphanumeric': 0.1234511243689766, 'fraction_numerical': 0.011748508490133088, 'mean_word_length': 3.315247524752475, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6168403', 'n_tokens_mistral': 4449, 'n_tokens_neox': 4173, 'n_words': 1125}