<commit_before>package neural

import (
	"fmt"
	"math/rand"
)

type TrainExample struct {
	Input  []float64
	Output []float64
}

// Evaluator wraps main tasks of NN, evaluate input data
type Evaluator interface {
	Evaluate(input []float64) []float64
	Train(trainExamples []TrainExample, epochs int, miniBatchSize int, learningRate float64)
}

type network struct {
	activator Activator
	layers    []Layer
}

func (n *network) Evaluate(input []float64) []float64 {
	output := input

	for _, layer := range n.layers {
		potentials := layer.Forward(output)
		output = n.Activate(potentials, true)
	}

	return output
}

func (n *network) Train(trainExamples []TrainExample, epochs int, miniBatchSize int, learningRate float64) {
	type Range struct {
		from, to int
	}

	samples := len(trainExamples)
	batches := samples / miniBatchSize
	if len(trainExamples)%miniBatchSize != 0 {
		batches++
	}

	batchRanges := make([]Range, batches, batches)
	for b := range batchRanges {
		min := b * miniBatchSize
		max := min + miniBatchSize
		if max > samples {
			max = samples
		}
		batchRanges[b] = Range{min, max}
	}

	for epoch := 0; epoch <= epochs; epoch++ {
		// Shuffle training data
		for i := range trainExamples {
			j := rand.Intn(i + 1)
			trainExamples[i], trainExamples[j] = trainExamples[j], trainExamples[i]
		}

		for b, batch := range batchRanges {
			t0 := time.Now()
			n.updateMiniBatch(trainExamples[batch.from:batch.to], learningRate)
			dt := time.Since(t0)
			if b%10 == 0 {
				fmt.Printf("%v/%v %v/%v    %v\r", epoch, epochs, b+1, batches, dt)
			}
		}
	}
}

type cn struct {
	biases  [][]float64
	weights [][][]float64
}

func (n *network) updateMiniBatch(miniBatch []TrainExample, learningRate float64) {
	layersCount := len(n.layers)
	samples := len(miniBatch)

	sumDeltaBias := make([][]float64, layersCount, layersCount)
	sumDeltaWeights := make([][][]float64, layersCount, layersCount)
	// sumDeltaBias := make([]*mat64.Dense, layersCount, layersCount)
	// sumDeltaWeights := make([]*mat64.Dense, layersCount, layersCount)
	buff := make(chan cn)

	for _, sample := range miniBatch {
		go func(sample TrainExample) {
			deltaWeights, deltaBias := n.backPropagation(sample)
			buff <- cn{biases: deltaBias, weights: deltaWeights}
		}(sample)
	}

	idx := 0
	for c := range buff {
		for l := range n.layers {
			if sumDeltaWeights[l] == nil {
				sumDeltaWeights[l] = copyOfMatrix(c.weights[l])
				// sumDeltaWeights[l] = mat64.DenseCopyOf(c.weights[l])
			} else {
				sumMatrix(sumDeltaWeights[l], c.weights[l])
				// sumDeltaWeights[l].Add(sumDeltaWeights[l], c.weights[l])
			}

			if sumDeltaBias[l] == nil {
				sumDeltaBias[l] = copyOfVector(c.biases[l])
				// sumDeltaBias[l] = mat64.DenseCopyOf(c.biases[l])
			} else {
				sumVector(sumDeltaBias[l], c.biases[l])
				// sumDeltaBias[l].Add(sumDeltaBias[l], c.biases[l])
			}
		}
		idx++
		if idx == samples {
			close(buff)
		}
	}

	rate := learningRate / float64(samples)
	for l, layer := range n.layers {
		// 	sumDeltaWeights[l].Scale(rate, sumDeltaWeights[l])
		// 	sumDeltaBias[l].Scale(rate, sumDeltaBias[l])
		// 	// fmt.Println(layer)
		mulVectorByScalar(sumDeltaBias[l], rate)
		mulMatrixByScalar(sumDeltaWeights[l], rate)
		layer.UpdateWeights(sumDeltaWeights[l], sumDeltaBias[l])
	}
}

func (n *network) backPropagation(sample TrainExample) (deltaWeights [][][]float64, deltaBias [][]float64) {
	layersCount := len(n.layers)

	acticationPerLayer := [][]float64{}
	potentialsPerLayer := [][]float64{}

	deltaBias = make([][]float64, layersCount, layersCount)
	deltaWeights = make([][][]float64, layersCount, layersCount)

	input := sample.Input
	acticationPerLayer = append(acticationPerLayer, input)
	for _, layer := range n.layers {
		potentials := layer.Forward(input)
		input = n.Activate(potentials, true)
		acticationPerLayer = append(acticationPerLayer, input)
		potentialsPerLayer = append(potentialsPerLayer, potentials)
	}

	errors := n.Diff(acticationPerLayer[len(acticationPerLayer)-1], sample.Output)
	delta := n.Delta(potentialsPerLayer[len(potentialsPerLayer)-1], errors)
	deltaBias[layersCount-1] = copyOfVector(delta)
	deltaWeights[layersCount-1] = mulTransposeVector(delta, acticationPerLayer[len(acticationPerLayer)-2])
	// deltaBias[layersCount-1] = mat64.NewDense(len(delta), 1, delta)
	// deltaWeights[layersCount-1] = n.MulTranspose(delta, acticationPerLayer[len(acticationPerLayer)-2])

	for l := 2; l <= layersCount; l++ {
		sp := n.Activate(potentialsPerLayer[len(potentialsPerLayer)-l], false)
		delta = n.Mul(n.layers[layersCount-l+1].Backward(delta), sp)
		deltaBias[layersCount-l] = copyOfVector(delta) // full copy can be avoided?
		deltaWeights[layersCount-l] = mulTransposeVector(delta, acticationPerLayer[len(acticationPerLayer)-l-1])
		// deltaBias[layersCount-l] = mat64.NewDense(len(delta), 1, delta)
		// deltaWeights[layersCount-l] = n.MulTranspose(delta, acticationPerLayer[len(acticationPerLayer)-l-1])
	}
	return
}

func (n *network) Activate(potentials []float64, forward bool) (output []float64) {
	output = make([]float64, len(potentials), len(potentials))

	if forward {
		for i, potential := range potentials {
			output[i] = n.activator.Activation(potential)
		}
	} else {
		for i, potential := range potentials {
			output[i] = n.activator.Derivative(potential)
		}
	}
	return
}

// func (n *network) MulTranspose(a, b []float64) (mul *mat64.Dense) {
// 	matA := mat64.NewDense(len(a), 1, a)
// 	matB := mat64.NewDense(len(b), 1, b)
// 	mul = mat64.NewDense(len(a), len(b), nil)

// 	mul.Mul(matA, matB.T())
// 	return
// }

func (n *network) Mul(a, b []float64) (mul []float64) {
	if len(a) != len(b) {
		errMsg := fmt.Sprintf("Incompatible sizes. %v vs %v", len(a), len(b))
		panic(errMsg)
	}

	mul = make([]float64, len(a), len(a))
	for i := range mul {
		mul[i] = a[i] * b[i]
	}
	return mul
}

func (n *network) Diff(a, b []float64) (diff []float64) {
	if len(a) != len(b) {
		errMsg := fmt.Sprintf("Incompatible sizes. %v vs %v", len(a), len(b))
		panic(errMsg)
	}

	diff = make([]float64, len(a), len(a))
	for i := range diff {
		diff[i] = a[i] - b[i]
	}
	return
}

func (n *network) Delta(potentials, errors []float64) (delta []float64) {
	if len(potentials) != len(errors) {
		errMsg := fmt.Sprintf("Incompatible sizes. %v vs %v", len(potentials), len(errors))
		panic(errMsg)
	}

	delta = make([]float64, len(potentials), len(potentials))
	for i := range potentials {
		delta[i] = errors[i] * n.activator.Derivative(potentials[i])
	}
	return
}

// NewNeuralNetwork initializes empty neural network
func NewNeuralNetwork(activator Activator, layers ...Layer) Evaluator {
	return &network{
		activator: activator,
		layers:    layers,
	}
}
<commit_msg>Add missing import<commit_after>package neural

import (
	"fmt"
	"math/rand"
	"time"
)

type TrainExample struct {
	Input  []float64
	Output []float64
}

// Evaluator wraps main tasks of NN, evaluate input data
type Evaluator interface {
	Evaluate(input []float64) []float64
	Train(trainExamples []TrainExample, epochs int, miniBatchSize int, learningRate float64)
}

type network struct {
	activator Activator
	layers    []Layer
}

func (n *network) Evaluate(input []float64) []float64 {
	output := input

	for _, layer := range n.layers {
		potentials := layer.Forward(output)
		output = n.Activate(potentials, true)
	}

	return output
}

func (n *network) Train(trainExamples []TrainExample, epochs int, miniBatchSize int, learningRate float64) {
	type Range struct {
		from, to int
	}

	samples := len(trainExamples)
	batches := samples / miniBatchSize
	if len(trainExamples)%miniBatchSize != 0 {
		batches++
	}

	batchRanges := make([]Range, batches, batches)
	for b := range batchRanges {
		min := b * miniBatchSize
		max := min + miniBatchSize
		if max > samples {
			max = samples
		}
		batchRanges[b] = Range{min, max}
	}

	for epoch := 0; epoch <= epochs; epoch++ {
		// Shuffle training data
		for i := range trainExamples {
			j := rand.Intn(i + 1)
			trainExamples[i], trainExamples[j] = trainExamples[j], trainExamples[i]
		}

		for b, batch := range batchRanges {
			t0 := time.Now()
			n.updateMiniBatch(trainExamples[batch.from:batch.to], learningRate)
			dt := time.Since(t0)
			if b%10 == 0 {
				fmt.Printf("%v/%v %v/%v    %v\r", epoch, epochs, b+1, batches, dt)
			}
		}
	}
}

type cn struct {
	biases  [][]float64
	weights [][][]float64
}

func (n *network) updateMiniBatch(miniBatch []TrainExample, learningRate float64) {
	layersCount := len(n.layers)
	samples := len(miniBatch)

	sumDeltaBias := make([][]float64, layersCount, layersCount)
	sumDeltaWeights := make([][][]float64, layersCount, layersCount)
	// sumDeltaBias := make([]*mat64.Dense, layersCount, layersCount)
	// sumDeltaWeights := make([]*mat64.Dense, layersCount, layersCount)
	buff := make(chan cn)

	for _, sample := range miniBatch {
		go func(sample TrainExample) {
			deltaWeights, deltaBias := n.backPropagation(sample)
			buff <- cn{biases: deltaBias, weights: deltaWeights}
		}(sample)
	}

	idx := 0
	for c := range buff {
		for l := range n.layers {
			if sumDeltaWeights[l] == nil {
				sumDeltaWeights[l] = copyOfMatrix(c.weights[l])
				// sumDeltaWeights[l] = mat64.DenseCopyOf(c.weights[l])
			} else {
				sumMatrix(sumDeltaWeights[l], c.weights[l])
				// sumDeltaWeights[l].Add(sumDeltaWeights[l], c.weights[l])
			}

			if sumDeltaBias[l] == nil {
				sumDeltaBias[l] = copyOfVector(c.biases[l])
				// sumDeltaBias[l] = mat64.DenseCopyOf(c.biases[l])
			} else {
				sumVector(sumDeltaBias[l], c.biases[l])
				// sumDeltaBias[l].Add(sumDeltaBias[l], c.biases[l])
			}
		}
		idx++
		if idx == samples {
			close(buff)
		}
	}

	rate := learningRate / float64(samples)
	for l, layer := range n.layers {
		// 	sumDeltaWeights[l].Scale(rate, sumDeltaWeights[l])
		// 	sumDeltaBias[l].Scale(rate, sumDeltaBias[l])
		// 	// fmt.Println(layer)
		mulVectorByScalar(sumDeltaBias[l], rate)
		mulMatrixByScalar(sumDeltaWeights[l], rate)
		layer.UpdateWeights(sumDeltaWeights[l], sumDeltaBias[l])
	}
}

func (n *network) backPropagation(sample TrainExample) (deltaWeights [][][]float64, deltaBias [][]float64) {
	layersCount := len(n.layers)

	acticationPerLayer := [][]float64{}
	potentialsPerLayer := [][]float64{}

	deltaBias = make([][]float64, layersCount, layersCount)
	deltaWeights = make([][][]float64, layersCount, layersCount)

	input := sample.Input
	acticationPerLayer = append(acticationPerLayer, input)
	for _, layer := range n.layers {
		potentials := layer.Forward(input)
		input = n.Activate(potentials, true)
		acticationPerLayer = append(acticationPerLayer, input)
		potentialsPerLayer = append(potentialsPerLayer, potentials)
	}

	errors := n.Diff(acticationPerLayer[len(acticationPerLayer)-1], sample.Output)
	delta := n.Delta(potentialsPerLayer[len(potentialsPerLayer)-1], errors)
	deltaBias[layersCount-1] = copyOfVector(delta)
	deltaWeights[layersCount-1] = mulTransposeVector(delta, acticationPerLayer[len(acticationPerLayer)-2])
	// deltaBias[layersCount-1] = mat64.NewDense(len(delta), 1, delta)
	// deltaWeights[layersCount-1] = n.MulTranspose(delta, acticationPerLayer[len(acticationPerLayer)-2])

	for l := 2; l <= layersCount; l++ {
		sp := n.Activate(potentialsPerLayer[len(potentialsPerLayer)-l], false)
		delta = n.Mul(n.layers[layersCount-l+1].Backward(delta), sp)
		deltaBias[layersCount-l] = copyOfVector(delta) // full copy can be avoided?
		deltaWeights[layersCount-l] = mulTransposeVector(delta, acticationPerLayer[len(acticationPerLayer)-l-1])
		// deltaBias[layersCount-l] = mat64.NewDense(len(delta), 1, delta)
		// deltaWeights[layersCount-l] = n.MulTranspose(delta, acticationPerLayer[len(acticationPerLayer)-l-1])
	}
	return
}

func (n *network) Activate(potentials []float64, forward bool) (output []float64) {
	output = make([]float64, len(potentials), len(potentials))

	if forward {
		for i, potential := range potentials {
			output[i] = n.activator.Activation(potential)
		}
	} else {
		for i, potential := range potentials {
			output[i] = n.activator.Derivative(potential)
		}
	}
	return
}

// func (n *network) MulTranspose(a, b []float64) (mul *mat64.Dense) {
// 	matA := mat64.NewDense(len(a), 1, a)
// 	matB := mat64.NewDense(len(b), 1, b)
// 	mul = mat64.NewDense(len(a), len(b), nil)

// 	mul.Mul(matA, matB.T())
// 	return
// }

func (n *network) Mul(a, b []float64) (mul []float64) {
	if len(a) != len(b) {
		errMsg := fmt.Sprintf("Incompatible sizes. %v vs %v", len(a), len(b))
		panic(errMsg)
	}

	mul = make([]float64, len(a), len(a))
	for i := range mul {
		mul[i] = a[i] * b[i]
	}
	return mul
}

func (n *network) Diff(a, b []float64) (diff []float64) {
	if len(a) != len(b) {
		errMsg := fmt.Sprintf("Incompatible sizes. %v vs %v", len(a), len(b))
		panic(errMsg)
	}

	diff = make([]float64, len(a), len(a))
	for i := range diff {
		diff[i] = a[i] - b[i]
	}
	return
}

func (n *network) Delta(potentials, errors []float64) (delta []float64) {
	if len(potentials) != len(errors) {
		errMsg := fmt.Sprintf("Incompatible sizes. %v vs %v", len(potentials), len(errors))
		panic(errMsg)
	}

	delta = make([]float64, len(potentials), len(potentials))
	for i := range potentials {
		delta[i] = errors[i] * n.activator.Derivative(potentials[i])
	}
	return
}

// NewNeuralNetwork initializes empty neural network
func NewNeuralNetwork(activator Activator, layers ...Layer) Evaluator {
	return &network{
		activator: activator,
		layers:    layers,
	}
}
<|endoftext|>


<commit_before>package google

import (
	"fmt"
	"log"
	"time"

	"google.golang.org/api/googleapi"
	"google.golang.org/api/compute/v1"

	"github.com/hashicorp/terraform/helper/hashcode"
	"github.com/hashicorp/terraform/helper/schema"
)

func resourceComputeInstanceGroupManager() *schema.Resource {
	return &schema.Resource{
		Create: resourceComputeInstanceGroupManagerCreate,
		Read:   resourceComputeInstanceGroupManagerRead,
		Update: resourceComputeInstanceGroupManagerUpdate,
		Delete: resourceComputeInstanceGroupManagerDelete,

		Schema: map[string]*schema.Schema{
			"name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"description": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
			},

			"base_instance_name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"fingerprint": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},

			"instance_group": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},

			"instance_template": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},

			"target_pools": &schema.Schema{
				Type:     schema.TypeSet,
				Required: true,
				Elem:     &schema.Schema{Type: schema.TypeString},
				Set: func(v interface{}) int {
					return hashcode.String(v.(string))
				},
			},

			"target_size": &schema.Schema{
				Type:     schema.TypeInt,
				Computed: true,
				Optional: true,
			},

			"zone": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"self_link": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
		},
	}
}

func waitOpZone(config *Config, op *compute.Operation, zone string,
	resource string, action string) (*compute.Operation, error) {

	w := &OperationWaiter{
		Service: config.clientCompute,
		Op:      op,
		Project: config.Project,
		Zone:    zone,
		Type:    OperationWaitZone,
	}
	state := w.Conf()
	state.Timeout = 2 * time.Minute
	state.MinTimeout = 1 * time.Second
	opRaw, err := state.WaitForState()
	if err != nil {
		return nil, fmt.Errorf("Error waiting for %s to %s: %s", resource, action, err)
	}
	return opRaw.(*compute.Operation), nil
}

func resourceComputeInstanceGroupManagerCreate(d *schema.ResourceData, meta interface{}) error {
	config := meta.(*Config)

	// Get group size, default to 1 if not given
	var target_size int64 = 1
	if v, ok := d.GetOk("target_size"); ok {
		target_size = int64(v.(int))
	}

	// Build the parameter
	manager := &compute.InstanceGroupManager{
		Name:             d.Get("name").(string),
		BaseInstanceName: d.Get("base_instance_name").(string),
		InstanceTemplate: d.Get("instance_template").(string),
		TargetSize: target_size,
	}

	// Set optional fields
	if v, ok := d.GetOk("description"); ok {
		manager.Description = v.(string)
	}

	if attr := d.Get("target_pools").(*schema.Set); attr.Len() > 0 {
		var s []string
		for _, v := range attr.List() {
			s = append(s, v.(string))
		}
		manager.TargetPools = s
	}

	log.Printf("[DEBUG] InstanceGroupManager insert request: %#v", manager)
	op, err := config.clientCompute.InstanceGroupManagers.Insert(
		config.Project, d.Get("zone").(string), manager).Do()
	if err != nil {
		return fmt.Errorf("Error creating InstanceGroupManager: %s", err)
	}

	// It probably maybe worked, so store the ID now
	d.SetId(manager.Name)

	// Wait for the operation to complete
	op, err = waitOpZone(config, op, d.Get("zone").(string), "InstanceGroupManager", "create")
	if err != nil {
		return err
	}
	if op.Error != nil {
		// The resource didn't actually create
		d.SetId("")
		// Return the error
		return OperationError(*op.Error)
	}

	return resourceComputeInstanceGroupManagerRead(d, meta)
}

func resourceComputeInstanceGroupManagerRead(d *schema.ResourceData, meta interface{}) error {
	config := meta.(*Config)

	manager, err := config.clientCompute.InstanceGroupManagers.Get(
		config.Project, d.Get("zone").(string), d.Id()).Do()
	if err != nil {
		if gerr, ok := err.(*googleapi.Error); ok && gerr.Code == 404 {
			// The resource doesn't exist anymore
			d.SetId("")

			return nil
		}

		return fmt.Errorf("Error reading instance group manager: %s", err)
	}

	// Set computed fields
	d.Set("fingerprint", manager.Fingerprint)
	d.Set("instance_group", manager.InstanceGroup)
	d.Set("target_size", manager.TargetSize)
	d.Set("self_link", manager.SelfLink)

	return nil
}
func resourceComputeInstanceGroupManagerUpdate(d *schema.ResourceData, meta interface{}) error {
	config := meta.(*Config)

	d.Partial(true)

	// If target_pools changes then update
	if d.HasChange("target_pools") {
		var targetPools []string
		if attr := d.Get("target_pools").(*schema.Set); attr.Len() > 0 {
			for _, v := range attr.List() {
				targetPools = append(targetPools, v.(string))
			}
		}

		// Build the parameter
		setTargetPools := &compute.InstanceGroupManagersSetTargetPoolsRequest{
			Fingerprint: d.Get("fingerprint").(string),
			TargetPools: targetPools,
		}

		op, err := config.clientCompute.InstanceGroupManagers.SetTargetPools(
			config.Project, d.Get("zone").(string), d.Id(), setTargetPools).Do()
		if err != nil {
			return fmt.Errorf("Error updating InstanceGroupManager: %s", err)
		}

		// Wait for the operation to complete
		op, err = waitOpZone(config, op, d.Get("zone").(string), "InstanceGroupManager", "update TargetPools")
		if err != nil {
			return err
		}
		if op.Error != nil {
			return OperationError(*op.Error)
		}

		d.SetPartial("target_pools")
	}

	// If instance_template changes then update
	if d.HasChange("instance_template") {
		// Build the parameter
		setInstanceTemplate := &compute.InstanceGroupManagersSetInstanceTemplateRequest{
			InstanceTemplate: d.Get("instance_template").(string),
		}

		op, err := config.clientCompute.InstanceGroupManagers.SetInstanceTemplate(
			config.Project, d.Get("zone").(string), d.Id(), setInstanceTemplate).Do()
		if err != nil {
			return fmt.Errorf("Error updating InstanceGroupManager: %s", err)
		}

		// Wait for the operation to complete
		op, err = waitOpZone(config, op, d.Get("zone").(string), "InstanceGroupManager", "update instance template")
		if err != nil {
			return err
		}
		if op.Error != nil {
			return OperationError(*op.Error)
		}

		d.SetPartial("instance_template")
	}

	// If size changes trigger a resize
	if d.HasChange("target_size") {
		if v, ok := d.GetOk("target_size"); ok {
			// Only do anything if the new size is set
			target_size := int64(v.(int))

			op, err := config.clientCompute.InstanceGroupManagers.Resize(
				config.Project, d.Get("zone").(string), d.Id(), target_size).Do()
			if err != nil {
				return fmt.Errorf("Error updating InstanceGroupManager: %s", err)
			}

			// Wait for the operation to complete
			op, err = waitOpZone(config, op, d.Get("zone").(string), "InstanceGroupManager", "update target_size")
			if err != nil {
				return err
			}
			if op.Error != nil {
				return OperationError(*op.Error)
			}
		}

		d.SetPartial("target_size")
	}

	d.Partial(false)

	return resourceComputeInstanceGroupManagerRead(d, meta)
}

func resourceComputeInstanceGroupManagerDelete(d *schema.ResourceData, meta interface{}) error {
	config := meta.(*Config)

	zone := d.Get("zone").(string)
	op, err := config.clientCompute.InstanceGroupManagers.Delete(config.Project, zone, d.Id()).Do()
	if err != nil {
		return fmt.Errorf("Error deleting instance group manager: %s", err)
	}

	// Wait for the operation to complete
	op, err = waitOpZone(config, op, d.Get("zone").(string), "InstanceGroupManager", "delete")
	if err != nil {
		return err
	}
	if op.Error != nil {
		// The resource didn't actually create
		d.SetId("")

		// Return the error
		return OperationError(*op.Error)
	}

	d.SetId("")
	return nil
}
<commit_msg>Increase timeout, IGM delete can be slow<commit_after>package google

import (
	"fmt"
	"log"
	"time"

	"google.golang.org/api/googleapi"
	"google.golang.org/api/compute/v1"

	"github.com/hashicorp/terraform/helper/hashcode"
	"github.com/hashicorp/terraform/helper/schema"
)

func resourceComputeInstanceGroupManager() *schema.Resource {
	return &schema.Resource{
		Create: resourceComputeInstanceGroupManagerCreate,
		Read:   resourceComputeInstanceGroupManagerRead,
		Update: resourceComputeInstanceGroupManagerUpdate,
		Delete: resourceComputeInstanceGroupManagerDelete,

		Schema: map[string]*schema.Schema{
			"name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"description": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
			},

			"base_instance_name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"fingerprint": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},

			"instance_group": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},

			"instance_template": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},

			"target_pools": &schema.Schema{
				Type:     schema.TypeSet,
				Required: true,
				Elem:     &schema.Schema{Type: schema.TypeString},
				Set: func(v interface{}) int {
					return hashcode.String(v.(string))
				},
			},

			"target_size": &schema.Schema{
				Type:     schema.TypeInt,
				Computed: true,
				Optional: true,
			},

			"zone": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"self_link": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
		},
	}
}

func waitOpZone(config *Config, op *compute.Operation, zone string,
	resource string, action string) (*compute.Operation, error) {

	w := &OperationWaiter{
		Service: config.clientCompute,
		Op:      op,
		Project: config.Project,
		Zone:    zone,
		Type:    OperationWaitZone,
	}
	state := w.Conf()
	state.Timeout = 8 * time.Minute
	state.MinTimeout = 1 * time.Second
	opRaw, err := state.WaitForState()
	if err != nil {
		return nil, fmt.Errorf("Error waiting for %s to %s: %s", resource, action, err)
	}
	return opRaw.(*compute.Operation), nil
}

func resourceComputeInstanceGroupManagerCreate(d *schema.ResourceData, meta interface{}) error {
	config := meta.(*Config)

	// Get group size, default to 1 if not given
	var target_size int64 = 1
	if v, ok := d.GetOk("target_size"); ok {
		target_size = int64(v.(int))
	}

	// Build the parameter
	manager := &compute.InstanceGroupManager{
		Name:             d.Get("name").(string),
		BaseInstanceName: d.Get("base_instance_name").(string),
		InstanceTemplate: d.Get("instance_template").(string),
		TargetSize: target_size,
	}

	// Set optional fields
	if v, ok := d.GetOk("description"); ok {
		manager.Description = v.(string)
	}

	if attr := d.Get("target_pools").(*schema.Set); attr.Len() > 0 {
		var s []string
		for _, v := range attr.List() {
			s = append(s, v.(string))
		}
		manager.TargetPools = s
	}

	log.Printf("[DEBUG] InstanceGroupManager insert request: %#v", manager)
	op, err := config.clientCompute.InstanceGroupManagers.Insert(
		config.Project, d.Get("zone").(string), manager).Do()
	if err != nil {
		return fmt.Errorf("Error creating InstanceGroupManager: %s", err)
	}

	// It probably maybe worked, so store the ID now
	d.SetId(manager.Name)

	// Wait for the operation to complete
	op, err = waitOpZone(config, op, d.Get("zone").(string), "InstanceGroupManager", "create")
	if err != nil {
		return err
	}
	if op.Error != nil {
		// The resource didn't actually create
		d.SetId("")
		// Return the error
		return OperationError(*op.Error)
	}

	return resourceComputeInstanceGroupManagerRead(d, meta)
}

func resourceComputeInstanceGroupManagerRead(d *schema.ResourceData, meta interface{}) error {
	config := meta.(*Config)

	manager, err := config.clientCompute.InstanceGroupManagers.Get(
		config.Project, d.Get("zone").(string), d.Id()).Do()
	if err != nil {
		if gerr, ok := err.(*googleapi.Error); ok && gerr.Code == 404 {
			// The resource doesn't exist anymore
			d.SetId("")

			return nil
		}

		return fmt.Errorf("Error reading instance group manager: %s", err)
	}

	// Set computed fields
	d.Set("fingerprint", manager.Fingerprint)
	d.Set("instance_group", manager.InstanceGroup)
	d.Set("target_size", manager.TargetSize)
	d.Set("self_link", manager.SelfLink)

	return nil
}
func resourceComputeInstanceGroupManagerUpdate(d *schema.ResourceData, meta interface{}) error {
	config := meta.(*Config)

	d.Partial(true)

	// If target_pools changes then update
	if d.HasChange("target_pools") {
		var targetPools []string
		if attr := d.Get("target_pools").(*schema.Set); attr.Len() > 0 {
			for _, v := range attr.List() {
				targetPools = append(targetPools, v.(string))
			}
		}

		// Build the parameter
		setTargetPools := &compute.InstanceGroupManagersSetTargetPoolsRequest{
			Fingerprint: d.Get("fingerprint").(string),
			TargetPools: targetPools,
		}

		op, err := config.clientCompute.InstanceGroupManagers.SetTargetPools(
			config.Project, d.Get("zone").(string), d.Id(), setTargetPools).Do()
		if err != nil {
			return fmt.Errorf("Error updating InstanceGroupManager: %s", err)
		}

		// Wait for the operation to complete
		op, err = waitOpZone(config, op, d.Get("zone").(string), "InstanceGroupManager", "update TargetPools")
		if err != nil {
			return err
		}
		if op.Error != nil {
			return OperationError(*op.Error)
		}

		d.SetPartial("target_pools")
	}

	// If instance_template changes then update
	if d.HasChange("instance_template") {
		// Build the parameter
		setInstanceTemplate := &compute.InstanceGroupManagersSetInstanceTemplateRequest{
			InstanceTemplate: d.Get("instance_template").(string),
		}

		op, err := config.clientCompute.InstanceGroupManagers.SetInstanceTemplate(
			config.Project, d.Get("zone").(string), d.Id(), setInstanceTemplate).Do()
		if err != nil {
			return fmt.Errorf("Error updating InstanceGroupManager: %s", err)
		}

		// Wait for the operation to complete
		op, err = waitOpZone(config, op, d.Get("zone").(string), "InstanceGroupManager", "update instance template")
		if err != nil {
			return err
		}
		if op.Error != nil {
			return OperationError(*op.Error)
		}

		d.SetPartial("instance_template")
	}

	// If size changes trigger a resize
	if d.HasChange("target_size") {
		if v, ok := d.GetOk("target_size"); ok {
			// Only do anything if the new size is set
			target_size := int64(v.(int))

			op, err := config.clientCompute.InstanceGroupManagers.Resize(
				config.Project, d.Get("zone").(string), d.Id(), target_size).Do()
			if err != nil {
				return fmt.Errorf("Error updating InstanceGroupManager: %s", err)
			}

			// Wait for the operation to complete
			op, err = waitOpZone(config, op, d.Get("zone").(string), "InstanceGroupManager", "update target_size")
			if err != nil {
				return err
			}
			if op.Error != nil {
				return OperationError(*op.Error)
			}
		}

		d.SetPartial("target_size")
	}

	d.Partial(false)

	return resourceComputeInstanceGroupManagerRead(d, meta)
}

func resourceComputeInstanceGroupManagerDelete(d *schema.ResourceData, meta interface{}) error {
	config := meta.(*Config)

	zone := d.Get("zone").(string)
	op, err := config.clientCompute.InstanceGroupManagers.Delete(config.Project, zone, d.Id()).Do()
	if err != nil {
		return fmt.Errorf("Error deleting instance group manager: %s", err)
	}

	// Wait for the operation to complete
	op, err = waitOpZone(config, op, d.Get("zone").(string), "InstanceGroupManager", "delete")
	if err != nil {
		return err
	}
	if op.Error != nil {
		// The resource didn't actually create
		d.SetId("")

		// Return the error
		return OperationError(*op.Error)
	}

	d.SetId("")
	return nil
}
<|endoftext|>


<commit_before>//  Copyright (c) 2019 Cisco and/or its affiliates.
//
//  Licensed 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.

package commands

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"path"
	"strings"
	"text/tabwriter"

	"github.com/sirupsen/logrus"
	"github.com/spf13/cobra"

	"go.ligato.io/vpp-agent/v2/cmd/agentctl/api/types"
	agentcli "go.ligato.io/vpp-agent/v2/cmd/agentctl/cli"
)

func NewModelCommand(cli agentcli.Cli) *cobra.Command {
	cmd := &cobra.Command{
		Use:   "model",
		Short: "Manage known models",
	}
	cmd.AddCommand(
		newModelListCommand(cli),
		newModelInspectCommand(cli),
	)
	return cmd
}

func newModelListCommand(cli agentcli.Cli) *cobra.Command {
	var opts ModelListOptions
	cmd := &cobra.Command{
		Use:     "ls [PATTERN]",
		Aliases: []string{"list", "l"},
		Short:   "List models",
		Args:    cobra.ArbitraryArgs,
		RunE: func(cmd *cobra.Command, args []string) error {
			opts.Refs = args
			return runModelList(cli, opts)
		},
	}
	flags := cmd.Flags()
	flags.StringVar(&opts.Class, "class", "", "Filter by model class")
	flags.StringVarP(&opts.Format, "format", "f", "", "Format output")
	return cmd
}

type ModelListOptions struct {
	Class  string
	Refs   []string
	Format string
}

func runModelList(cli agentcli.Cli, opts ModelListOptions) error {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	allModels, err := cli.Client().ModelList(ctx, types.ModelListOptions{
		Class: opts.Class,
	})
	if err != nil {
		return err
	}

	models := filterModelsByRefs(allModels, opts.Refs)

	format := opts.Format
	if len(format) == 0 {
		printModelTable(cli.Out(), models)
	} else {
		if err := formatAsTemplate(cli.Out(), format, models); err != nil {
			return err
		}
	}

	return nil
}

func printModelTable(out io.Writer, models []types.Model) {
	var buf bytes.Buffer
	w := tabwriter.NewWriter(&buf, 0, 0, 2, ' ', 0)
	fmt.Fprintf(w, "MODEL\tCLASS\tPROTO MESSAGE\tKEY PREFIX\t\n")
	for _, model := range models {
		fmt.Fprintf(w, "%s\t%s\t%s\t%s\t\n",
			model.Name, model.Class, model.ProtoName, model.KeyPrefix)
	}
	if err := w.Flush(); err != nil {
		panic(err)
	}

	fmt.Fprint(out, buf.String())
}

func filterModelsByPrefix(models []types.Model, prefixes []string) ([]types.Model, error) {
	if len(prefixes) == 0 {
		return models, nil
	}
	var filtered []types.Model
	for _, pref := range prefixes {
		var model types.Model
		for _, m := range models {
			if !strings.HasPrefix(m.Name, pref) {
				continue
			}
			if model.Name != "" {
				return nil, fmt.Errorf("multiple models found with provided prefix: %s", pref)
			}
			model = m
		}
		if model.Name == "" {
			return nil, fmt.Errorf("no model found for provided prefix: %s", pref)
		}
		filtered = append(filtered, model)
	}
	return filtered, nil
}

func filterModelsByRefs(models []types.Model, refs []string) []types.Model {
	var filtered []types.Model
	for _, model := range models {
		if !matchAnyRef(model, refs) {
			continue
		}
		filtered = append(filtered, model)
	}
	return filtered
}

func matchAnyRef(model types.Model, refs []string) bool {
	if len(refs) == 0 {
		return true
	}
	for _, ref := range refs {
		if ok, _ := path.Match(ref, model.Name); ok {
			return true
		}
	}
	return false
}

func newModelInspectCommand(cli agentcli.Cli) *cobra.Command {
	var (
		opts ModelInspectOptions
	)
	cmd := &cobra.Command{
		Use:     "inspect MODEL [MODEL...]",
		Aliases: []string{"i"},
		Short:   "Display detailed information on one or more models",
		Args:    cobra.MinimumNArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			opts.Names = args
			return runModelInspect(cli, opts)
		},
	}
	// TODO: add support for custom formatting instead of json
	//cmd.Flags().StringVar(&opts.Format, "format", "", "Format for the output")
	return cmd
}

type ModelInspectOptions struct {
	Names  []string
	Format string
}

func runModelInspect(cli agentcli.Cli, opts ModelInspectOptions) error {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	allModels, err := cli.Client().ModelList(ctx, types.ModelListOptions{})
	if err != nil {
		return err
	}

	models, err := filterModelsByPrefix(allModels, opts.Names)
	if err != nil {
		return err
	}

	logrus.Debugf("models: %+v", models)

	b, err := json.MarshalIndent(models, "", "  ")
	if err != nil {
		return fmt.Errorf("encoding data failed: %v", err)
	}

	fmt.Fprintf(cli.Out(), "%s\n", b)
	return nil
}
<commit_msg>Support custom format for model inspect<commit_after>//  Copyright (c) 2019 Cisco and/or its affiliates.
//
//  Licensed 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.

package commands

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"path"
	"strings"
	"text/tabwriter"

	"github.com/sirupsen/logrus"
	"github.com/spf13/cobra"

	"go.ligato.io/vpp-agent/v2/cmd/agentctl/api/types"
	agentcli "go.ligato.io/vpp-agent/v2/cmd/agentctl/cli"
)

func NewModelCommand(cli agentcli.Cli) *cobra.Command {
	cmd := &cobra.Command{
		Use:   "model",
		Short: "Manage known models",
	}
	cmd.AddCommand(
		newModelListCommand(cli),
		newModelInspectCommand(cli),
	)
	return cmd
}

func newModelListCommand(cli agentcli.Cli) *cobra.Command {
	var opts ModelListOptions
	cmd := &cobra.Command{
		Use:     "ls [PATTERN]",
		Aliases: []string{"list", "l"},
		Short:   "List models",
		Args:    cobra.ArbitraryArgs,
		RunE: func(cmd *cobra.Command, args []string) error {
			opts.Refs = args
			return runModelList(cli, opts)
		},
	}
	flags := cmd.Flags()
	flags.StringVar(&opts.Class, "class", "", "Filter by model class")
	flags.StringVarP(&opts.Format, "format", "f", "", "Format output")
	return cmd
}

type ModelListOptions struct {
	Class  string
	Refs   []string
	Format string
}

func runModelList(cli agentcli.Cli, opts ModelListOptions) error {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	allModels, err := cli.Client().ModelList(ctx, types.ModelListOptions{
		Class: opts.Class,
	})
	if err != nil {
		return err
	}

	models := filterModelsByRefs(allModels, opts.Refs)

	format := opts.Format
	if len(format) == 0 {
		printModelTable(cli.Out(), models)
	} else {
		if err := formatAsTemplate(cli.Out(), format, models); err != nil {
			return err
		}
	}

	return nil
}

func printModelTable(out io.Writer, models []types.Model) {
	var buf bytes.Buffer
	w := tabwriter.NewWriter(&buf, 0, 0, 2, ' ', 0)
	fmt.Fprintf(w, "MODEL\tCLASS\tPROTO MESSAGE\tKEY PREFIX\t\n")
	for _, model := range models {
		fmt.Fprintf(w, "%s\t%s\t%s\t%s\t\n",
			model.Name, model.Class, model.ProtoName, model.KeyPrefix)
	}
	if err := w.Flush(); err != nil {
		panic(err)
	}

	fmt.Fprint(out, buf.String())
}

func filterModelsByPrefix(models []types.Model, prefixes []string) ([]types.Model, error) {
	if len(prefixes) == 0 {
		return models, nil
	}
	var filtered []types.Model
	for _, pref := range prefixes {
		var model types.Model
		for _, m := range models {
			if !strings.HasPrefix(m.Name, pref) {
				continue
			}
			if model.Name != "" {
				return nil, fmt.Errorf("multiple models found with provided prefix: %s", pref)
			}
			model = m
		}
		if model.Name == "" {
			return nil, fmt.Errorf("no model found for provided prefix: %s", pref)
		}
		filtered = append(filtered, model)
	}
	return filtered, nil
}

func filterModelsByRefs(models []types.Model, refs []string) []types.Model {
	var filtered []types.Model
	for _, model := range models {
		if !matchAnyRef(model, refs) {
			continue
		}
		filtered = append(filtered, model)
	}
	return filtered
}

func matchAnyRef(model types.Model, refs []string) bool {
	if len(refs) == 0 {
		return true
	}
	for _, ref := range refs {
		if ok, _ := path.Match(ref, model.Name); ok {
			return true
		}
	}
	return false
}

func newModelInspectCommand(cli agentcli.Cli) *cobra.Command {
	var (
		opts ModelInspectOptions
	)
	cmd := &cobra.Command{
		Use:     "inspect MODEL [MODEL...]",
		Aliases: []string{"i"},
		Short:   "Display detailed information on one or more models",
		Args:    cobra.MinimumNArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			opts.Names = args
			return runModelInspect(cli, opts)
		},
	}
	cmd.Flags().StringVarP(&opts.Format, "format", "f", "", "Format for the output")
	return cmd
}

type ModelInspectOptions struct {
	Names  []string
	Format string
}

func runModelInspect(cli agentcli.Cli, opts ModelInspectOptions) error {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	allModels, err := cli.Client().ModelList(ctx, types.ModelListOptions{})
	if err != nil {
		return err
	}

	models, err := filterModelsByPrefix(allModels, opts.Names)
	if err != nil {
		return err
	}

	logrus.Debugf("models: %+v", models)

	format := opts.Format
	if len(format) == 0 {
		format = "json"
	}

	if err := formatAsTemplate(cli.Out(), format, models); err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2018 Adam Shannon
//
// Licensed 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.

package certutil

import (
	"crypto/x509/pkix"
	"encoding/asn1"
	"fmt"
	"regexp"
	"strings"
)

var (
	oidCommonName = asn1.ObjectIdentifier{2, 5, 4, 3}
)

// TODO(adam): Replace with RDSSquence.String() in g1.10
func StringifyPKIXName(name pkix.Name) (out string) {
	if len(name.OrganizationalUnit) > 0 {
		out = fmt.Sprintf("%s, %s", strings.Join(name.Organization, " "), name.OrganizationalUnit[0])
	}

	if out == "" {
		out = strings.Join(name.Organization, " ")
	}

	for i := range name.Names {
		if name.Names[i].Type.Equal(oidCommonName) {
			s, ok := name.Names[i].Value.(string)
			if ok {
				return cleanPKIXName(s)
			}
		}
	}

	return cleanPKIXName(out)
}

// Remove annoying characters from PKIX names
// e.g. newlines, line feeds, tabs, etc
func cleanPKIXName(name string) string {
	space := " "
	stripper := strings.NewReplacer("\n", space, "\r\n", space, "\t", space, "\r", space, "\f", space, "\v", space)
	name = stripper.Replace(name)

	trimmer := regexp.MustCompile(`(\s{1,})`)
	return trimmer.ReplaceAllString(name, " ")
}
<commit_msg>certutik/pkix: drop TODO for pkix.Name.String()<commit_after>// Copyright 2018 Adam Shannon
//
// Licensed 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.

package certutil

import (
	"crypto/x509/pkix"
	"encoding/asn1"
	"fmt"
	"regexp"
	"strings"
)

var (
	oidCommonName = asn1.ObjectIdentifier{2, 5, 4, 3}
)

func StringifyPKIXName(name pkix.Name) (out string) {
	if len(name.OrganizationalUnit) > 0 {
		out = fmt.Sprintf("%s, %s", strings.Join(name.Organization, " "), name.OrganizationalUnit[0])
	}

	if out == "" {
		out = strings.Join(name.Organization, " ")
	}

	for i := range name.Names {
		if name.Names[i].Type.Equal(oidCommonName) {
			s, ok := name.Names[i].Value.(string)
			if ok {
				return cleanPKIXName(s)
			}
		}
	}

	return cleanPKIXName(out)
}

// Remove annoying characters from PKIX names
// e.g. newlines, line feeds, tabs, etc
func cleanPKIXName(name string) string {
	space := " "
	stripper := strings.NewReplacer("\n", space, "\r\n", space, "\t", space, "\r", space, "\f", space, "\v", space)
	name = stripper.Replace(name)

	trimmer := regexp.MustCompile(`(\s{1,})`)
	return trimmer.ReplaceAllString(name, " ")
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

Licensed 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.
*/

package e2e_node

import (
	"os/exec"
	"path/filepath"
	"regexp"
	"time"

	"k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/uuid"
	"k8s.io/kubernetes/pkg/features"
	"k8s.io/kubernetes/pkg/kubelet/apis/kubeletconfig"
	"k8s.io/kubernetes/test/e2e/framework"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

const (
	devicePluginFeatureGate = "DevicePlugins=true"
	testPodNamePrefix       = "nvidia-gpu-"
)

// Serial because the test restarts Kubelet
var _ = framework.KubeDescribe("NVIDIA GPU Device Plugin [Feature:GPUDevicePlugin] [Serial] [Disruptive]", func() {
	f := framework.NewDefaultFramework("device-plugin-gpus-errors")

	Context("DevicePlugin", func() {
		By("Enabling support for Device Plugin")
		tempSetCurrentKubeletConfig(f, func(initialConfig *kubeletconfig.KubeletConfiguration) {
			initialConfig.FeatureGates[string(features.DevicePlugins)] = true
		})

		BeforeEach(func() {
			By("Ensuring that Nvidia GPUs exists on the node")
			if !checkIfNvidiaGPUsExistOnNode() {
				Skip("Nvidia GPUs do not exist on the node. Skipping test.")
			}

			By("Creating the Google Device Plugin pod for NVIDIA GPU in GKE")
			f.PodClient().CreateSync(framework.NVIDIADevicePlugin(f.Namespace.Name))

			By("Waiting for GPUs to become available on the local node")
			Eventually(func() bool {
				return framework.NumberOfNVIDIAGPUs(getLocalNode(f)) > 0
			}, 10*time.Second, framework.Poll).Should(BeTrue())

			if framework.NumberOfNVIDIAGPUs(getLocalNode(f)) < 2 {
				Skip("Not enough GPUs to execute this test (at least two needed)")
			}
		})

		AfterEach(func() {
			l, err := f.PodClient().List(metav1.ListOptions{})
			framework.ExpectNoError(err)

			for _, p := range l.Items {
				if p.Namespace != f.Namespace.Name {
					continue
				}

				f.PodClient().Delete(p.Name, &metav1.DeleteOptions{})
			}
		})

		It("checks that when Kubelet restarts exclusive GPU assignation to pods is kept.", func() {
			By("Creating one GPU pod on a node with at least two GPUs")
			p1 := f.PodClient().CreateSync(makeCudaPauseImage())
			devId1 := getDeviceId(f, p1.Name, p1.Name, 1)
			p1, err := f.PodClient().Get(p1.Name, metav1.GetOptions{})
			framework.ExpectNoError(err)

			By("Restarting Kubelet and waiting for the current running pod to restart")
			restartKubelet(f)

			By("Confirming that after a kubelet and pod restart, GPU assignement is kept")
			devIdRestart := getDeviceId(f, p1.Name, p1.Name, 2)
			Expect(devIdRestart).To(Equal(devId1))

			By("Restarting Kubelet and creating another pod")
			restartKubelet(f)
			p2 := f.PodClient().CreateSync(makeCudaPauseImage())

			By("Checking that pods got a different GPU")
			devId2 := getDeviceId(f, p2.Name, p2.Name, 1)
			Expect(devId1).To(Not(Equal(devId2)))

			// Cleanup
			f.PodClient().DeleteSync(p1.Name, &metav1.DeleteOptions{}, framework.DefaultPodDeletionTimeout)
			f.PodClient().DeleteSync(p2.Name, &metav1.DeleteOptions{}, framework.DefaultPodDeletionTimeout)
		})
	})
})

func makeCudaPauseImage() *v1.Pod {
	podName := testPodNamePrefix + string(uuid.NewUUID())

	return &v1.Pod{
		ObjectMeta: metav1.ObjectMeta{Name: podName},
		Spec: v1.PodSpec{
			RestartPolicy: v1.RestartPolicyAlways,
			Containers: []v1.Container{{
				Image: busyboxImage,
				Name:  podName,
				// Retrieves the gpu devices created in the user pod.
				// Note the nvidia device plugin implementation doesn't do device id remapping currently.
				// Will probably need to use nvidia-smi if that changes.
				Command: []string{"sh", "-c", "devs=$(ls /dev/ | egrep '^nvidia[0-9]+$') && echo gpu devices: $devs"},

				Resources: v1.ResourceRequirements{
					Limits:   newDecimalResourceList(framework.NVIDIAGPUResourceName, 1),
					Requests: newDecimalResourceList(framework.NVIDIAGPUResourceName, 1),
				},
			}},
		},
	}
}

func newDecimalResourceList(name v1.ResourceName, quantity int64) v1.ResourceList {
	return v1.ResourceList{name: *resource.NewQuantity(quantity, resource.DecimalSI)}
}

// TODO: Find a uniform way to deal with systemctl/initctl/service operations. #34494
func restartKubelet(f *framework.Framework) {
	beforeSocks, err := filepath.Glob("/var/lib/kubelet/device-plugins/nvidiaGPU*.sock")
	framework.ExpectNoError(err)
	Expect(len(beforeSocks)).NotTo(BeZero())
	stdout, err := exec.Command("sudo", "systemctl", "list-units", "kubelet*", "--state=running").CombinedOutput()
	framework.ExpectNoError(err)
	regex := regexp.MustCompile("(kubelet-[0-9]+)")
	matches := regex.FindStringSubmatch(string(stdout))
	Expect(len(matches)).NotTo(BeZero())
	kube := matches[0]
	framework.Logf("Get running kubelet with systemctl: %v, %v", string(stdout), kube)
	stdout, err = exec.Command("sudo", "systemctl", "restart", kube).CombinedOutput()
	framework.ExpectNoError(err, "Failed to restart kubelet with systemctl: %v, %v", err, stdout)
	Eventually(func() ([]string, error) {
		return filepath.Glob("/var/lib/kubelet/device-plugins/nvidiaGPU*.sock")
	}, 5*time.Minute, framework.Poll).ShouldNot(ConsistOf(beforeSocks))
}

func getDeviceId(f *framework.Framework, podName string, contName string, restartCount int32) string {
	// Wait till pod has been restarted at least restartCount times.
	Eventually(func() bool {
		p, err := f.PodClient().Get(podName, metav1.GetOptions{})
		if err != nil || len(p.Status.ContainerStatuses) < 1 {
			return false
		}
		return p.Status.ContainerStatuses[0].RestartCount >= restartCount
	}, 5*time.Minute, framework.Poll).Should(BeTrue())
	logs, err := framework.GetPodLogs(f.ClientSet, f.Namespace.Name, podName, contName)
	if err != nil {
		framework.Failf("GetPodLogs for pod %q failed: %v", podName, err)
	}
	framework.Logf("got pod logs: %v", logs)
	regex := regexp.MustCompile("gpu devices: (nvidia[0-9]+)")
	matches := regex.FindStringSubmatch(logs)
	if len(matches) < 2 {
		return ""
	}
	return matches[1]
}
<commit_msg>Extends gpu_device_plugin e2e_node test to verify that scheduled pods can continue to run even after device plugin deletion and kubelet restarts.<commit_after>/*
Copyright 2017 The Kubernetes Authors.

Licensed 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.
*/

package e2e_node

import (
	"os/exec"
	"regexp"
	"time"

	"k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/uuid"
	"k8s.io/kubernetes/pkg/features"
	"k8s.io/kubernetes/pkg/kubelet/apis/kubeletconfig"
	"k8s.io/kubernetes/test/e2e/framework"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

const (
	devicePluginFeatureGate = "DevicePlugins=true"
	testPodNamePrefix       = "nvidia-gpu-"
)

// Serial because the test restarts Kubelet
var _ = framework.KubeDescribe("NVIDIA GPU Device Plugin [Feature:GPUDevicePlugin] [Serial] [Disruptive]", func() {
	f := framework.NewDefaultFramework("device-plugin-gpus-errors")

	Context("DevicePlugin", func() {
		By("Enabling support for Device Plugin")
		tempSetCurrentKubeletConfig(f, func(initialConfig *kubeletconfig.KubeletConfiguration) {
			initialConfig.FeatureGates[string(features.DevicePlugins)] = true
		})

		var devicePluginPod *v1.Pod
		BeforeEach(func() {
			By("Ensuring that Nvidia GPUs exists on the node")
			if !checkIfNvidiaGPUsExistOnNode() {
				Skip("Nvidia GPUs do not exist on the node. Skipping test.")
			}

			By("Creating the Google Device Plugin pod for NVIDIA GPU in GKE")
			devicePluginPod = f.PodClient().CreateSync(framework.NVIDIADevicePlugin(f.Namespace.Name))

			By("Waiting for GPUs to become available on the local node")
			Eventually(func() bool {
				return framework.NumberOfNVIDIAGPUs(getLocalNode(f)) > 0
			}, 10*time.Second, framework.Poll).Should(BeTrue())

			if framework.NumberOfNVIDIAGPUs(getLocalNode(f)) < 2 {
				Skip("Not enough GPUs to execute this test (at least two needed)")
			}
		})

		AfterEach(func() {
			l, err := f.PodClient().List(metav1.ListOptions{})
			framework.ExpectNoError(err)

			for _, p := range l.Items {
				if p.Namespace != f.Namespace.Name {
					continue
				}

				f.PodClient().Delete(p.Name, &metav1.DeleteOptions{})
			}
		})

		It("checks that when Kubelet restarts exclusive GPU assignation to pods is kept.", func() {
			By("Creating one GPU pod on a node with at least two GPUs")
			p1 := f.PodClient().CreateSync(makeCudaPauseImage())
			count1, devId1 := getDeviceId(f, p1.Name, p1.Name, 1)
			p1, err := f.PodClient().Get(p1.Name, metav1.GetOptions{})
			framework.ExpectNoError(err)

			By("Restarting Kubelet and waiting for the current running pod to restart")
			restartKubelet(f)

			By("Confirming that after a kubelet and pod restart, GPU assignement is kept")
			count1, devIdRestart1 := getDeviceId(f, p1.Name, p1.Name, count1+1)
			Expect(devIdRestart1).To(Equal(devId1))

			By("Restarting Kubelet and creating another pod")
			restartKubelet(f)
			p2 := f.PodClient().CreateSync(makeCudaPauseImage())

			By("Checking that pods got a different GPU")
			count2, devId2 := getDeviceId(f, p2.Name, p2.Name, 1)
			Expect(devId1).To(Not(Equal(devId2)))

			By("Deleting device plugin.")
			f.PodClient().Delete(devicePluginPod.Name, &metav1.DeleteOptions{})
			By("Waiting for GPUs to become unavailable on the local node")
			Eventually(func() bool {
				return framework.NumberOfNVIDIAGPUs(getLocalNode(f)) <= 0
			}, 10*time.Minute, framework.Poll).Should(BeTrue())
			By("Checking that scheduled pods can continue to run even after we delete device plugin.")
			count1, devIdRestart1 = getDeviceId(f, p1.Name, p1.Name, count1+1)
			Expect(devIdRestart1).To(Equal(devId1))
			count2, devIdRestart2 := getDeviceId(f, p2.Name, p2.Name, count2+1)
			Expect(devIdRestart2).To(Equal(devId2))
			By("Restarting Kubelet.")
			restartKubelet(f)
			By("Checking that scheduled pods can continue to run even after we delete device plugin and restart Kubelet.")
			count1, devIdRestart1 = getDeviceId(f, p1.Name, p1.Name, count1+2)
			Expect(devIdRestart1).To(Equal(devId1))
			count2, devIdRestart2 = getDeviceId(f, p2.Name, p2.Name, count2+2)
			Expect(devIdRestart2).To(Equal(devId2))

			// Cleanup
			f.PodClient().DeleteSync(p1.Name, &metav1.DeleteOptions{}, framework.DefaultPodDeletionTimeout)
			f.PodClient().DeleteSync(p2.Name, &metav1.DeleteOptions{}, framework.DefaultPodDeletionTimeout)
		})
	})
})

func makeCudaPauseImage() *v1.Pod {
	podName := testPodNamePrefix + string(uuid.NewUUID())

	return &v1.Pod{
		ObjectMeta: metav1.ObjectMeta{Name: podName},
		Spec: v1.PodSpec{
			RestartPolicy: v1.RestartPolicyAlways,
			Containers: []v1.Container{{
				Image: busyboxImage,
				Name:  podName,
				// Retrieves the gpu devices created in the user pod.
				// Note the nvidia device plugin implementation doesn't do device id remapping currently.
				// Will probably need to use nvidia-smi if that changes.
				Command: []string{"sh", "-c", "devs=$(ls /dev/ | egrep '^nvidia[0-9]+$') && echo gpu devices: $devs"},

				Resources: v1.ResourceRequirements{
					Limits:   newDecimalResourceList(framework.NVIDIAGPUResourceName, 1),
					Requests: newDecimalResourceList(framework.NVIDIAGPUResourceName, 1),
				},
			}},
		},
	}
}

func newDecimalResourceList(name v1.ResourceName, quantity int64) v1.ResourceList {
	return v1.ResourceList{name: *resource.NewQuantity(quantity, resource.DecimalSI)}
}

// TODO: Find a uniform way to deal with systemctl/initctl/service operations. #34494
func restartKubelet(f *framework.Framework) {
	stdout, err := exec.Command("sudo", "systemctl", "list-units", "kubelet*", "--state=running").CombinedOutput()
	framework.ExpectNoError(err)
	regex := regexp.MustCompile("(kubelet-[0-9]+)")
	matches := regex.FindStringSubmatch(string(stdout))
	Expect(len(matches)).NotTo(BeZero())
	kube := matches[0]
	framework.Logf("Get running kubelet with systemctl: %v, %v", string(stdout), kube)
	stdout, err = exec.Command("sudo", "systemctl", "restart", kube).CombinedOutput()
	framework.ExpectNoError(err, "Failed to restart kubelet with systemctl: %v, %v", err, stdout)
}

func getDeviceId(f *framework.Framework, podName string, contName string, restartCount int32) (int32, string) {
	var count int32
	// Wait till pod has been restarted at least restartCount times.
	Eventually(func() bool {
		p, err := f.PodClient().Get(podName, metav1.GetOptions{})
		if err != nil || len(p.Status.ContainerStatuses) < 1 {
			return false
		}
		count = p.Status.ContainerStatuses[0].RestartCount
		return count >= restartCount
	}, 5*time.Minute, framework.Poll).Should(BeTrue())
	logs, err := framework.GetPodLogs(f.ClientSet, f.Namespace.Name, podName, contName)
	if err != nil {
		framework.Failf("GetPodLogs for pod %q failed: %v", podName, err)
	}
	framework.Logf("got pod logs: %v", logs)
	regex := regexp.MustCompile("gpu devices: (nvidia[0-9]+)")
	matches := regex.FindStringSubmatch(logs)
	if len(matches) < 2 {
		return count, ""
	}
	return count, matches[1]
}
<|endoftext|>


<commit_before>package neurgo

import (
	"encoding/json"
	"errors"
	"fmt"
	"github.com/proxypoke/vector"
	"log"
	"sync"
)

type Neuron struct {
	NodeId             *NodeId
	Bias               float64
	Inbound            []*InboundConnection
	Outbound           []*OutboundConnection
	Closing            chan chan bool
	DataChan           chan *DataMessage
	ActivationFunction *EncodableActivation
	wg                 *sync.WaitGroup
	Cortex             *Cortex
}

func (neuron *Neuron) MarshalJSON() ([]byte, error) {
	return json.Marshal(
		struct {
			NodeId             *NodeId
			Bias               float64
			Inbound            []*InboundConnection
			Outbound           []*OutboundConnection
			ActivationFunction *EncodableActivation
		}{
			NodeId:             neuron.NodeId,
			Bias:               neuron.Bias,
			Inbound:            neuron.Inbound,
			Outbound:           neuron.Outbound,
			ActivationFunction: neuron.ActivationFunction,
		})
}

func (neuron *Neuron) Run() {

	log.Printf("%v Run() started", neuron.NodeId.UUID)

	defer neuron.wg.Done()

	neuron.checkRunnable()

	neuron.sendEmptySignalRecurrentOutbound()

	weightedInputs := createEmptyWeightedInputs(neuron.Inbound)

	closed := false

	for {

		log.Printf("Neuron %v select().  datachan: %v", neuron.NodeId.UUID, neuron.DataChan)

		select {
		case responseChan := <-neuron.Closing:
			closed = true
			responseChan <- true
			break // TODO: do we need this for anything??
		case dataMessage := <-neuron.DataChan:
			log.Printf("Neuron %v recording input: %v", neuron.NodeId.UUID, dataMessage)
			recordInput(weightedInputs, dataMessage)
			log.Printf("Neuron %v new weightedInputs: %v", neuron.NodeId.UUID, weightedInputs)
		}

		if closed {
			neuron.Closing = nil
			neuron.DataChan = nil
			break
		}

		if receiveBarrierSatisfied(weightedInputs) {

			log.Printf("Neuron %v barrier satisfied via inputs: %v", neuron.NodeId.UUID, weightedInputs)
			scalarOutput := neuron.computeScalarOutput(weightedInputs)

			dataMessage := &DataMessage{
				SenderId: neuron.NodeId,
				Inputs:   []float64{scalarOutput},
			}

			neuron.scatterOutput(dataMessage)

			weightedInputs = createEmptyWeightedInputs(neuron.Inbound)

		} else {
			log.Printf("Neuron %v receive barrier not satisfied.  weightedInputs: %v", neuron.NodeId.UUID, weightedInputs)
		}

	}

	log.Printf("%v Run() finished", neuron.NodeId.UUID)

}

func (neuron *Neuron) String() string {
	return JsonString(neuron)
}

func (neuron *Neuron) ConnectOutbound(connectable OutboundConnectable) {
	if neuron.Outbound == nil {
		neuron.Outbound = make([]*OutboundConnection, 0)
	}
	connection := &OutboundConnection{
		NodeId:   connectable.nodeId(),
		DataChan: connectable.dataChan(),
	}
	neuron.Outbound = append(neuron.Outbound, connection)
}

func (neuron *Neuron) ConnectInboundWeighted(connectable InboundConnectable, weights []float64) *InboundConnection {
	if neuron.Inbound == nil {
		neuron.Inbound = make([]*InboundConnection, 0)
	}
	connection := &InboundConnection{
		NodeId:  connectable.nodeId(),
		Weights: weights,
	}
	neuron.Inbound = append(neuron.Inbound, connection)
	return connection
}

// In order to prevent deadlock, any neurons we have recurrent outbound
// connections to must be "primed" by sending an empty signal.  A recurrent
// outbound connection simply means that it's a connection to ourself or
// to a neuron in a previous (eg, to the left) layer.  If we didn't do this,
// that previous neuron would be waiting forever for a signal that will
// never come, because this neuron wouldn't fire until it got a signal.
func (neuron *Neuron) sendEmptySignalRecurrentOutbound() {

	recurrentConnections := neuron.recurrentOutboundConnections()
	for _, recurrentConnection := range recurrentConnections {

		inputs := []float64{0}
		dataMessage := &DataMessage{
			SenderId: neuron.NodeId,
			Inputs:   inputs,
		}
		recurrentConnection.DataChan <- dataMessage
	}

}

// Find the subset of outbound connections which are "recurrent" - meaning
// that the connection is to this neuron itself, or to a neuron in a previous
// (eg, to the left) layer.
func (neuron *Neuron) recurrentOutboundConnections() []*OutboundConnection {
	result := make([]*OutboundConnection, 0)
	for _, outboundConnection := range neuron.Outbound {
		if neuron.isConnectionRecurrent(outboundConnection) {
			result = append(result, outboundConnection)
		}
	}
	return result
}

// a connection is considered recurrent if it has a connection
// to itself or to a node in a previous layer.  Previous meaning
// if you look at a feedforward from left to right, with the input
// layer being on the far left, and output layer on the far right,
// then any layer to the left is considered previous.
func (neuron *Neuron) isConnectionRecurrent(connection *OutboundConnection) bool {
	if connection.NodeId.LayerIndex <= neuron.NodeId.LayerIndex {
		return true
	}
	return false
}

// same as isConnectionRecurrent, but for inbound connections
// TODO: use interfaces to eliminate code duplication
func (neuron *Neuron) IsInboundConnectionRecurrent(connection *InboundConnection) bool {
	if neuron.NodeId.LayerIndex <= connection.NodeId.LayerIndex {
		return true
	}
	return false
}

func (neuron *Neuron) scatterOutput(dataMessage *DataMessage) {
	for _, outboundConnection := range neuron.Outbound {
		dataChan := outboundConnection.DataChan
		log.Printf("Neuron %v scatter %v to: %v", neuron.NodeId.UUID, dataMessage, outboundConnection)
		dataChan <- dataMessage
	}
}

func (neuron *Neuron) Init() {
	if neuron.Closing == nil {
		neuron.Closing = make(chan chan bool)
	}

	if neuron.DataChan == nil {
		neuron.DataChan = make(chan *DataMessage, len(neuron.Inbound))
	}

	/*
		if neuron.ActivationFunction == nil {

			// TODO: fix this .. we need to serialize the name of
			// the function, and when we deserialize, resolve to
			// actual function
			neuron.ActivationFunction = EncodableSigmoid()
		}
	*/

	if neuron.wg == nil {
		neuron.wg = &sync.WaitGroup{}
		neuron.wg.Add(1)
	}

}

func (neuron *Neuron) Shutdown() {

	closingResponse := make(chan bool)
	neuron.Closing <- closingResponse
	response := <-closingResponse
	if response != true {
		log.Panicf("Got unexpected response on closing channel")
	}

	neuron.shutdownOutboundConnections()

	neuron.wg.Wait()
	neuron.wg = nil
}

func (neuron *Neuron) InboundUUIDMap() UUIDToInboundConnection {
	inboundUUIDMap := make(UUIDToInboundConnection)
	for _, connection := range neuron.Inbound {
		inboundUUIDMap[connection.NodeId.UUID] = connection
	}
	return inboundUUIDMap
}

func (neuron *Neuron) Copy() *Neuron {

	// serialize to json
	jsonBytes, err := json.Marshal(neuron)
	if err != nil {
		log.Fatal(err)
	}

	// new neuron
	neuronCopy := &Neuron{}

	// deserialize json into new neuron
	err = json.Unmarshal(jsonBytes, neuronCopy)
	if err != nil {
		log.Fatal(err)
	}

	return neuronCopy

}

func (neuron *Neuron) checkRunnable() {

	if neuron.NodeId == nil {
		msg := fmt.Sprintf("not expecting neuron.NodeId to be nil")
		panic(msg)
	}

	if neuron.Inbound == nil {
		msg := fmt.Sprintf("not expecting neuron.Inbound to be nil")
		panic(msg)
	}

	if neuron.Closing == nil {
		msg := fmt.Sprintf("not expecting neuron.Closing to be nil")
		panic(msg)
	}

	if neuron.DataChan == nil {
		msg := fmt.Sprintf("not expecting neuron.DataChan to be nil")
		panic(msg)
	}

	if neuron.ActivationFunction == nil {
		msg := fmt.Sprintf("not expecting neuron.ActivationFunction to be nil")
		panic(msg)
	}

	if err := neuron.validateOutbound(); err != nil {
		msg := fmt.Sprintf("invalid outbound connection(s): %v", err.Error())
		panic(msg)
	}

}

func (neuron *Neuron) validateOutbound() error {
	for _, connection := range neuron.Outbound {
		if connection.DataChan == nil {
			msg := fmt.Sprintf("%v has empty DataChan", connection)
			return errors.New(msg)
		}
	}
	return nil
}

func (neuron *Neuron) computeScalarOutput(weightedInputs []*weightedInput) float64 {
	output := neuron.weightedInputDotProductSum(weightedInputs)
	output += neuron.Bias
	output = neuron.ActivationFunction.ActivationFunction(output)
	return output
}

// for each weighted input vector, calculate the (inputs * weights) dot product
// and sum all of these dot products together to produce a sum
func (neuron *Neuron) weightedInputDotProductSum(weightedInputs []*weightedInput) float64 {

	var dotProductSummation float64
	dotProductSummation = 0

	for _, weightedInput := range weightedInputs {
		inputs := weightedInput.inputs
		weights := weightedInput.weights
		inputVector := vector.NewFrom(inputs)
		weightVector := vector.NewFrom(weights)
		log.Printf("inputVector: %v", inputVector)
		log.Printf("weightVector: %v", weightVector)
		dotProduct, error := vector.DotProduct(inputVector, weightVector)
		if error != nil {
			t := "%T error performing dot product between %v and %v"
			message := fmt.Sprintf(t, neuron, inputVector, weightVector)
			panic(message)
		}
		dotProductSummation += dotProduct
	}

	return dotProductSummation

}

func (neuron *Neuron) dataChan() chan *DataMessage {
	return neuron.DataChan
}

func (neuron *Neuron) nodeId() *NodeId {
	return neuron.NodeId
}

func (neuron *Neuron) initOutboundConnections(nodeIdToDataMsg nodeIdToDataMsgMap) {
	for _, outboundConnection := range neuron.Outbound {
		if outboundConnection.DataChan == nil {
			dataChan := nodeIdToDataMsg[outboundConnection.NodeId.UUID]
			if dataChan != nil {
				outboundConnection.DataChan = dataChan
			}
		}
	}
}

func (neuron *Neuron) shutdownOutboundConnections() {
	for _, outboundConnection := range neuron.Outbound {
		outboundConnection.DataChan = nil
	}
}
<commit_msg>return connection on function that makes new connection<commit_after>package neurgo

import (
	"encoding/json"
	"errors"
	"fmt"
	"github.com/proxypoke/vector"
	"log"
	"sync"
)

type Neuron struct {
	NodeId             *NodeId
	Bias               float64
	Inbound            []*InboundConnection
	Outbound           []*OutboundConnection
	Closing            chan chan bool
	DataChan           chan *DataMessage
	ActivationFunction *EncodableActivation
	wg                 *sync.WaitGroup
	Cortex             *Cortex
}

func (neuron *Neuron) MarshalJSON() ([]byte, error) {
	return json.Marshal(
		struct {
			NodeId             *NodeId
			Bias               float64
			Inbound            []*InboundConnection
			Outbound           []*OutboundConnection
			ActivationFunction *EncodableActivation
		}{
			NodeId:             neuron.NodeId,
			Bias:               neuron.Bias,
			Inbound:            neuron.Inbound,
			Outbound:           neuron.Outbound,
			ActivationFunction: neuron.ActivationFunction,
		})
}

func (neuron *Neuron) Run() {

	log.Printf("%v Run() started", neuron.NodeId.UUID)

	defer neuron.wg.Done()

	neuron.checkRunnable()

	neuron.sendEmptySignalRecurrentOutbound()

	weightedInputs := createEmptyWeightedInputs(neuron.Inbound)

	closed := false

	for {

		log.Printf("Neuron %v select().  datachan: %v", neuron.NodeId.UUID, neuron.DataChan)

		select {
		case responseChan := <-neuron.Closing:
			closed = true
			responseChan <- true
			break // TODO: do we need this for anything??
		case dataMessage := <-neuron.DataChan:
			log.Printf("Neuron %v recording input: %v", neuron.NodeId.UUID, dataMessage)
			recordInput(weightedInputs, dataMessage)
			log.Printf("Neuron %v new weightedInputs: %v", neuron.NodeId.UUID, weightedInputs)
		}

		if closed {
			neuron.Closing = nil
			neuron.DataChan = nil
			break
		}

		if receiveBarrierSatisfied(weightedInputs) {

			log.Printf("Neuron %v barrier satisfied via inputs: %v", neuron.NodeId.UUID, weightedInputs)
			scalarOutput := neuron.computeScalarOutput(weightedInputs)

			dataMessage := &DataMessage{
				SenderId: neuron.NodeId,
				Inputs:   []float64{scalarOutput},
			}

			neuron.scatterOutput(dataMessage)

			weightedInputs = createEmptyWeightedInputs(neuron.Inbound)

		} else {
			log.Printf("Neuron %v receive barrier not satisfied.  weightedInputs: %v", neuron.NodeId.UUID, weightedInputs)
		}

	}

	log.Printf("%v Run() finished", neuron.NodeId.UUID)

}

func (neuron *Neuron) String() string {
	return JsonString(neuron)
}

func (neuron *Neuron) ConnectOutbound(connectable OutboundConnectable) *OutboundConnection {
	if neuron.Outbound == nil {
		neuron.Outbound = make([]*OutboundConnection, 0)
	}
	connection := &OutboundConnection{
		NodeId:   connectable.nodeId(),
		DataChan: connectable.dataChan(),
	}
	neuron.Outbound = append(neuron.Outbound, connection)
	return connection
}

func (neuron *Neuron) ConnectInboundWeighted(connectable InboundConnectable, weights []float64) *InboundConnection {
	if neuron.Inbound == nil {
		neuron.Inbound = make([]*InboundConnection, 0)
	}
	connection := &InboundConnection{
		NodeId:  connectable.nodeId(),
		Weights: weights,
	}
	neuron.Inbound = append(neuron.Inbound, connection)
	return connection
}

// In order to prevent deadlock, any neurons we have recurrent outbound
// connections to must be "primed" by sending an empty signal.  A recurrent
// outbound connection simply means that it's a connection to ourself or
// to a neuron in a previous (eg, to the left) layer.  If we didn't do this,
// that previous neuron would be waiting forever for a signal that will
// never come, because this neuron wouldn't fire until it got a signal.
func (neuron *Neuron) sendEmptySignalRecurrentOutbound() {

	recurrentConnections := neuron.recurrentOutboundConnections()
	for _, recurrentConnection := range recurrentConnections {

		inputs := []float64{0}
		dataMessage := &DataMessage{
			SenderId: neuron.NodeId,
			Inputs:   inputs,
		}
		recurrentConnection.DataChan <- dataMessage
	}

}

// Find the subset of outbound connections which are "recurrent" - meaning
// that the connection is to this neuron itself, or to a neuron in a previous
// (eg, to the left) layer.
func (neuron *Neuron) recurrentOutboundConnections() []*OutboundConnection {
	result := make([]*OutboundConnection, 0)
	for _, outboundConnection := range neuron.Outbound {
		if neuron.isConnectionRecurrent(outboundConnection) {
			result = append(result, outboundConnection)
		}
	}
	return result
}

// a connection is considered recurrent if it has a connection
// to itself or to a node in a previous layer.  Previous meaning
// if you look at a feedforward from left to right, with the input
// layer being on the far left, and output layer on the far right,
// then any layer to the left is considered previous.
func (neuron *Neuron) isConnectionRecurrent(connection *OutboundConnection) bool {
	if connection.NodeId.LayerIndex <= neuron.NodeId.LayerIndex {
		return true
	}
	return false
}

// same as isConnectionRecurrent, but for inbound connections
// TODO: use interfaces to eliminate code duplication
func (neuron *Neuron) IsInboundConnectionRecurrent(connection *InboundConnection) bool {
	if neuron.NodeId.LayerIndex <= connection.NodeId.LayerIndex {
		return true
	}
	return false
}

func (neuron *Neuron) scatterOutput(dataMessage *DataMessage) {
	for _, outboundConnection := range neuron.Outbound {
		dataChan := outboundConnection.DataChan
		log.Printf("Neuron %v scatter %v to: %v", neuron.NodeId.UUID, dataMessage, outboundConnection)
		dataChan <- dataMessage
	}
}

func (neuron *Neuron) Init() {
	if neuron.Closing == nil {
		neuron.Closing = make(chan chan bool)
	}

	if neuron.DataChan == nil {
		neuron.DataChan = make(chan *DataMessage, len(neuron.Inbound))
	}

	/*
		if neuron.ActivationFunction == nil {

			// TODO: fix this .. we need to serialize the name of
			// the function, and when we deserialize, resolve to
			// actual function
			neuron.ActivationFunction = EncodableSigmoid()
		}
	*/

	if neuron.wg == nil {
		neuron.wg = &sync.WaitGroup{}
		neuron.wg.Add(1)
	}

}

func (neuron *Neuron) Shutdown() {

	closingResponse := make(chan bool)
	neuron.Closing <- closingResponse
	response := <-closingResponse
	if response != true {
		log.Panicf("Got unexpected response on closing channel")
	}

	neuron.shutdownOutboundConnections()

	neuron.wg.Wait()
	neuron.wg = nil
}

func (neuron *Neuron) InboundUUIDMap() UUIDToInboundConnection {
	inboundUUIDMap := make(UUIDToInboundConnection)
	for _, connection := range neuron.Inbound {
		inboundUUIDMap[connection.NodeId.UUID] = connection
	}
	return inboundUUIDMap
}

func (neuron *Neuron) Copy() *Neuron {

	// serialize to json
	jsonBytes, err := json.Marshal(neuron)
	if err != nil {
		log.Fatal(err)
	}

	// new neuron
	neuronCopy := &Neuron{}

	// deserialize json into new neuron
	err = json.Unmarshal(jsonBytes, neuronCopy)
	if err != nil {
		log.Fatal(err)
	}

	return neuronCopy

}

func (neuron *Neuron) checkRunnable() {

	if neuron.NodeId == nil {
		msg := fmt.Sprintf("not expecting neuron.NodeId to be nil")
		panic(msg)
	}

	if neuron.Inbound == nil {
		msg := fmt.Sprintf("not expecting neuron.Inbound to be nil")
		panic(msg)
	}

	if neuron.Closing == nil {
		msg := fmt.Sprintf("not expecting neuron.Closing to be nil")
		panic(msg)
	}

	if neuron.DataChan == nil {
		msg := fmt.Sprintf("not expecting neuron.DataChan to be nil")
		panic(msg)
	}

	if neuron.ActivationFunction == nil {
		msg := fmt.Sprintf("not expecting neuron.ActivationFunction to be nil")
		panic(msg)
	}

	if err := neuron.validateOutbound(); err != nil {
		msg := fmt.Sprintf("invalid outbound connection(s): %v", err.Error())
		panic(msg)
	}

}

func (neuron *Neuron) validateOutbound() error {
	for _, connection := range neuron.Outbound {
		if connection.DataChan == nil {
			msg := fmt.Sprintf("%v has empty DataChan", connection)
			return errors.New(msg)
		}
	}
	return nil
}

func (neuron *Neuron) computeScalarOutput(weightedInputs []*weightedInput) float64 {
	output := neuron.weightedInputDotProductSum(weightedInputs)
	output += neuron.Bias
	output = neuron.ActivationFunction.ActivationFunction(output)
	return output
}

// for each weighted input vector, calculate the (inputs * weights) dot product
// and sum all of these dot products together to produce a sum
func (neuron *Neuron) weightedInputDotProductSum(weightedInputs []*weightedInput) float64 {

	var dotProductSummation float64
	dotProductSummation = 0

	for _, weightedInput := range weightedInputs {
		inputs := weightedInput.inputs
		weights := weightedInput.weights
		inputVector := vector.NewFrom(inputs)
		weightVector := vector.NewFrom(weights)
		log.Printf("inputVector: %v", inputVector)
		log.Printf("weightVector: %v", weightVector)
		dotProduct, error := vector.DotProduct(inputVector, weightVector)
		if error != nil {
			t := "%T error performing dot product between %v and %v"
			message := fmt.Sprintf(t, neuron, inputVector, weightVector)
			panic(message)
		}
		dotProductSummation += dotProduct
	}

	return dotProductSummation

}

func (neuron *Neuron) dataChan() chan *DataMessage {
	return neuron.DataChan
}

func (neuron *Neuron) nodeId() *NodeId {
	return neuron.NodeId
}

func (neuron *Neuron) initOutboundConnections(nodeIdToDataMsg nodeIdToDataMsgMap) {
	for _, outboundConnection := range neuron.Outbound {
		if outboundConnection.DataChan == nil {
			dataChan := nodeIdToDataMsg[outboundConnection.NodeId.UUID]
			if dataChan != nil {
				outboundConnection.DataChan = dataChan
			}
		}
	}
}

func (neuron *Neuron) shutdownOutboundConnections() {
	for _, outboundConnection := range neuron.Outbound {
		outboundConnection.DataChan = nil
	}
}
<|endoftext|>


<commit_before>package postgres

import (
	"database/sql"
	"fmt"
	. "github.com/aktau/gomig/db/common"
	"log"
	"strings"
)

var PG_W_VERBOSE = true

var (
	postgresInit = []string{
		"SET client_encoding = 'UTF8'",
		"SET standard_conforming_strings = off",
		"SET check_function_bodies = false",
		"SET client_min_messages = warning",
	}
)

const (
	explainQuery = `
SELECT col.column_name AS field,
       CASE
        WHEN col.character_maximum_length IS NOT NULL THEN col.data_type || '(' || col.character_maximum_length || ')'
        ELSE col.data_type
       END AS type,
       col.is_nullable AS null,
       CASE
        WHEN tc.constraint_type = 'PRIMARY KEY' THEN 'PRI'
        ELSE ''
       END AS key,
       '' AS default,
       '' AS extra
       --kcu.constraint_name AS constraint_name
       --kcu.*,
       --tc.*
FROM   information_schema.columns col
LEFT JOIN   information_schema.key_column_usage kcu ON (kcu.table_name = col.table_name AND kcu.column_name = col.column_name)
LEFT JOIN   information_schema.table_constraints AS tc ON (kcu.constraint_name = tc.constraint_name)
WHERE  col.table_name = '%v'
ORDER BY col.ordinal_position;`
)

type genericPostgresWriter struct {
	e               Executor
	insertBulkLimit int
}

/* how to do an UPSERT/MERGE in PostgreSQL
 * http://stackoverflow.com/questions/17267417/how-do-i-do-an-upsert-merge-insert-on-duplicate-update-in-postgresq */
func (w *genericPostgresWriter) MergeTable(src *Table, dstName string, r Reader) error {
	tmpName := "gomig_tmp"
	stmts := make([]string, 0, 5)

	/* create temporary table */
	stmts = append(stmts,
		fmt.Sprintf("CREATE TEMPORARY TABLE %v (\n\t%v\n)\nON COMMIT DROP;\n", tmpName, ColumnsSql(src)))

	if PG_W_VERBOSE {
		log.Println("MergeTable: preparing to read values")
	}

	/* bulk insert values */
	rows, err := r.Read(src)
	if err != nil {
		return err
	}
	defer rows.Close()

	if PG_W_VERBOSE {
		log.Println("MergeTable: query done, scanning rows...")
	}

	/* an alternate way to do this, with type assertions
	 * but possibly less accurately: http://go-database-sql.org/varcols.html */
	pointers := make([]interface{}, len(src.Columns))
	containers := make([]sql.RawBytes, len(src.Columns))
	for i, _ := range pointers {
		pointers[i] = &containers[i]
	}
	stringrep := make([]string, 0, len(src.Columns))
	insertLines := make([]string, 0, 32)
	for rows.Next() {
		if PG_W_VERBOSE {
			log.Println("MergeTable: inside a loop, copying number of values:", len(src.Columns))
		}

		err := rows.Scan(pointers...)
		if err != nil {
			log.Println("MergeTable: error while reading from source:", err)
			return err
		}

		for idx, val := range containers {
			if val == nil {
				stringrep = append(stringrep, "NULL")
			} else {
				switch src.Columns[idx].Type {
				case "text":
					stringrep = append(stringrep, "$$"+string(val)+"$$")
				case "boolean":
					/* ascii(48) = "0" and ascii(49) = "1" */
					switch val[0] {
					case 48:
						stringrep = append(stringrep, "f")
					case 49:
						stringrep = append(stringrep, "t")
					default:
						return fmt.Errorf("writer: did not recognize bool value: string(%v) = %v, val[0] = %v", val, string(val), val[0])
					}
				case "integer":
					stringrep = append(stringrep, string(val))
				default:
					stringrep = append(stringrep, string(val))
				}
			}
		}

		insertLines = append(insertLines, "("+strings.Join(stringrep, ",")+")")
		stringrep = stringrep[:0]

		if len(insertLines) > w.insertBulkLimit {
			stmts = append(stmts, fmt.Sprintf("INSERT INTO %v VALUES\n\t%v;\n",
				tmpName, strings.Join(insertLines, "\n\t")))

			insertLines = insertLines[:0]
		}
	}

	err = rows.Err()
	if err != nil {
		return err
	}
	if len(insertLines) > 0 {
		stmts = append(stmts, fmt.Sprintf("INSERT INTO %v VALUES\n\t%v;\n",
			tmpName, strings.Join(insertLines, "\n\t")))
	}

	/* analyze the temp table, for performance */
	stmts = append(stmts, fmt.Sprintf("ANALYZE %v;\n", tmpName))

	/* lock the target table */
	stmts = append(stmts, fmt.Sprintf("LOCK TABLE %v IN EXCLUSIVE MODE;", dstName))

	colnames := make([]string, 0, len(src.Columns))
	srccol := make([]string, 0, len(src.Columns))
	pkWhere := make([]string, 0, len(src.Columns))
	pkIsNull := make([]string, 0, len(src.Columns))
	colassign := make([]string, 0, len(src.Columns))
	for _, col := range src.Columns {
		colnames = append(colnames, col.Name)
		srccol = append(srccol, "src."+col.Name)
		if col.PrimaryKey {
			pkWhere = append(pkWhere, fmt.Sprintf("dst.%[1]v = src.%[1]v", col.Name))
			pkIsNull = append(pkIsNull, fmt.Sprintf("dst.%[1]v IS NULL", col.Name))
		} else {
			colassign = append(colassign, fmt.Sprintf("dst.%[1]v = src.%[1]v", col.Name))
		}
	}
	pkWherePart := strings.Join(pkWhere, "\nAND    ")
	pkIsNullPart := strings.Join(pkIsNull, "\nAND    ")
	srccolPart := strings.Join(srccol, ",\n       ")

	/* UPDATE from temp table to target table based on PK */
	stmts = append(stmts, fmt.Sprintf(`
UPDATE %v AS dst
SET    %v
FROM   %v AS src
WHERE  %v;`, dstName, strings.Join(colassign, ",\n       "), tmpName, pkWherePart))

	/* INSERT from temp table to target table based on PK */
	stmts = append(stmts, fmt.Sprintf(`
INSERT INTO %[1]v (%[3]v)
SELECT %[4]v
FROM   %[2]v AS src
LEFT OUTER JOIN %[1]v AS dst ON (
	   %[5]v
)
WHERE  %[6]v;
`, dstName, tmpName, strings.Join(colnames, ", "), srccolPart, pkWherePart, pkIsNullPart))

	err = w.e.Transaction(
		fmt.Sprintf("merge table %v into table %v", src.Name, dstName), stmts)
	return err
}

func (w *genericPostgresWriter) Close() error {
	return w.e.Close()
}

type PostgresWriter struct {
	genericPostgresWriter
}

func NewPostgresWriter(conf *Config) (*PostgresWriter, error) {
	db, err := openDB(conf)
	if err != nil {
		return nil, err
	}

	executor, err := NewDbExecutor(db)
	if err != nil {
		db.Close()
		return nil, err
	}

	errors := executor.Multiple("initializing DB connection (WARNING: connection pooling might mess with this)", postgresInit)
	if len(errors) > 0 {
		executor.Close()
		for _, err := range errors {
			log.Println("postgres error:", err)
		}
		return nil, errors[0]
	}

	return &PostgresWriter{genericPostgresWriter{executor, 64}}, nil
}

type PostgresFileWriter struct {
	genericPostgresWriter
}

func NewPostgresFileWriter(filename string) (*PostgresFileWriter, error) {
	executor, err := NewFileExecutor(filename)
	if err != nil {
		return nil, err
	}

	errors := executor.Multiple("initializing DB connection", postgresInit)
	if len(errors) > 0 {
		executor.Close()
		for _, err := range errors {
			log.Println("postgres error:", err)
		}
		return nil, errors[0]
	}

	return &PostgresFileWriter{genericPostgresWriter{executor, 256}}, err
}

func PostgresType(genericType string) string {
	return genericType
}

func ColumnsSql(table *Table) string {
	colSql := make([]string, 0, len(table.Columns))

	for _, col := range table.Columns {
		colSql = append(colSql, fmt.Sprintf("%v %v", col.Name, PostgresType(col.Type)))
	}

	return strings.Join(colSql, ",\n\t")
}
<commit_msg>forgot comma for bulk insert<commit_after>package postgres

import (
	"database/sql"
	"fmt"
	. "github.com/aktau/gomig/db/common"
	"log"
	"strings"
)

var PG_W_VERBOSE = true

var (
	postgresInit = []string{
		"SET client_encoding = 'UTF8'",
		"SET standard_conforming_strings = off",
		"SET check_function_bodies = false",
		"SET client_min_messages = warning",
	}
)

const (
	explainQuery = `
SELECT col.column_name AS field,
       CASE
        WHEN col.character_maximum_length IS NOT NULL THEN col.data_type || '(' || col.character_maximum_length || ')'
        ELSE col.data_type
       END AS type,
       col.is_nullable AS null,
       CASE
        WHEN tc.constraint_type = 'PRIMARY KEY' THEN 'PRI'
        ELSE ''
       END AS key,
       '' AS default,
       '' AS extra
       --kcu.constraint_name AS constraint_name
       --kcu.*,
       --tc.*
FROM   information_schema.columns col
LEFT JOIN   information_schema.key_column_usage kcu ON (kcu.table_name = col.table_name AND kcu.column_name = col.column_name)
LEFT JOIN   information_schema.table_constraints AS tc ON (kcu.constraint_name = tc.constraint_name)
WHERE  col.table_name = '%v'
ORDER BY col.ordinal_position;`
)

type genericPostgresWriter struct {
	e               Executor
	insertBulkLimit int
}

/* how to do an UPSERT/MERGE in PostgreSQL
 * http://stackoverflow.com/questions/17267417/how-do-i-do-an-upsert-merge-insert-on-duplicate-update-in-postgresq */
func (w *genericPostgresWriter) MergeTable(src *Table, dstName string, r Reader) error {
	tmpName := "gomig_tmp"
	stmts := make([]string, 0, 5)

	/* create temporary table */
	stmts = append(stmts,
		fmt.Sprintf("CREATE TEMPORARY TABLE %v (\n\t%v\n)\nON COMMIT DROP;\n", tmpName, ColumnsSql(src)))

	if PG_W_VERBOSE {
		log.Println("MergeTable: preparing to read values")
	}

	/* bulk insert values */
	rows, err := r.Read(src)
	if err != nil {
		return err
	}
	defer rows.Close()

	if PG_W_VERBOSE {
		log.Println("MergeTable: query done, scanning rows...")
	}

	/* an alternate way to do this, with type assertions
	 * but possibly less accurately: http://go-database-sql.org/varcols.html */
	pointers := make([]interface{}, len(src.Columns))
	containers := make([]sql.RawBytes, len(src.Columns))
	for i, _ := range pointers {
		pointers[i] = &containers[i]
	}
	stringrep := make([]string, 0, len(src.Columns))
	insertLines := make([]string, 0, 32)
	for rows.Next() {
		if PG_W_VERBOSE {
			log.Println("MergeTable: inside a loop, copying number of values:", len(src.Columns))
		}

		err := rows.Scan(pointers...)
		if err != nil {
			log.Println("MergeTable: error while reading from source:", err)
			return err
		}

		for idx, val := range containers {
			if val == nil {
				stringrep = append(stringrep, "NULL")
			} else {
				switch src.Columns[idx].Type {
				case "text":
					stringrep = append(stringrep, "$$"+string(val)+"$$")
				case "boolean":
					/* ascii(48) = "0" and ascii(49) = "1" */
					switch val[0] {
					case 48:
						stringrep = append(stringrep, "f")
					case 49:
						stringrep = append(stringrep, "t")
					default:
						return fmt.Errorf("writer: did not recognize bool value: string(%v) = %v, val[0] = %v", val, string(val), val[0])
					}
				case "integer":
					stringrep = append(stringrep, string(val))
				default:
					stringrep = append(stringrep, string(val))
				}
			}
		}

		insertLines = append(insertLines, "("+strings.Join(stringrep, ",")+")")
		stringrep = stringrep[:0]

		if len(insertLines) > w.insertBulkLimit {
			stmts = append(stmts, fmt.Sprintf("INSERT INTO %v VALUES\n\t%v;\n",
				tmpName, strings.Join(insertLines, ",\n\t")))

			insertLines = insertLines[:0]
		}
	}

	err = rows.Err()
	if err != nil {
		return err
	}
	if len(insertLines) > 0 {
		stmts = append(stmts, fmt.Sprintf("INSERT INTO %v VALUES\n\t%v;\n",
			tmpName, strings.Join(insertLines, ",\n\t")))
	}

	/* analyze the temp table, for performance */
	stmts = append(stmts, fmt.Sprintf("ANALYZE %v;\n", tmpName))

	/* lock the target table */
	stmts = append(stmts, fmt.Sprintf("LOCK TABLE %v IN EXCLUSIVE MODE;", dstName))

	colnames := make([]string, 0, len(src.Columns))
	srccol := make([]string, 0, len(src.Columns))
	pkWhere := make([]string, 0, len(src.Columns))
	pkIsNull := make([]string, 0, len(src.Columns))
	colassign := make([]string, 0, len(src.Columns))
	for _, col := range src.Columns {
		colnames = append(colnames, col.Name)
		srccol = append(srccol, "src."+col.Name)
		if col.PrimaryKey {
			pkWhere = append(pkWhere, fmt.Sprintf("dst.%[1]v = src.%[1]v", col.Name))
			pkIsNull = append(pkIsNull, fmt.Sprintf("dst.%[1]v IS NULL", col.Name))
		} else {
			colassign = append(colassign, fmt.Sprintf("dst.%[1]v = src.%[1]v", col.Name))
		}
	}
	pkWherePart := strings.Join(pkWhere, "\nAND    ")
	pkIsNullPart := strings.Join(pkIsNull, "\nAND    ")
	srccolPart := strings.Join(srccol, ",\n       ")

	/* UPDATE from temp table to target table based on PK */
	stmts = append(stmts, fmt.Sprintf(`
UPDATE %v AS dst
SET    %v
FROM   %v AS src
WHERE  %v;`, dstName, strings.Join(colassign, ",\n       "), tmpName, pkWherePart))

	/* INSERT from temp table to target table based on PK */
	stmts = append(stmts, fmt.Sprintf(`
INSERT INTO %[1]v (%[3]v)
SELECT %[4]v
FROM   %[2]v AS src
LEFT OUTER JOIN %[1]v AS dst ON (
	   %[5]v
)
WHERE  %[6]v;
`, dstName, tmpName, strings.Join(colnames, ", "), srccolPart, pkWherePart, pkIsNullPart))

	err = w.e.Transaction(
		fmt.Sprintf("merge table %v into table %v", src.Name, dstName), stmts)
	return err
}

func (w *genericPostgresWriter) Close() error {
	return w.e.Close()
}

type PostgresWriter struct {
	genericPostgresWriter
}

func NewPostgresWriter(conf *Config) (*PostgresWriter, error) {
	db, err := openDB(conf)
	if err != nil {
		return nil, err
	}

	executor, err := NewDbExecutor(db)
	if err != nil {
		db.Close()
		return nil, err
	}

	errors := executor.Multiple("initializing DB connection (WARNING: connection pooling might mess with this)", postgresInit)
	if len(errors) > 0 {
		executor.Close()
		for _, err := range errors {
			log.Println("postgres error:", err)
		}
		return nil, errors[0]
	}

	return &PostgresWriter{genericPostgresWriter{executor, 64}}, nil
}

type PostgresFileWriter struct {
	genericPostgresWriter
}

func NewPostgresFileWriter(filename string) (*PostgresFileWriter, error) {
	executor, err := NewFileExecutor(filename)
	if err != nil {
		return nil, err
	}

	errors := executor.Multiple("initializing DB connection", postgresInit)
	if len(errors) > 0 {
		executor.Close()
		for _, err := range errors {
			log.Println("postgres error:", err)
		}
		return nil, errors[0]
	}

	return &PostgresFileWriter{genericPostgresWriter{executor, 256}}, err
}

func PostgresType(genericType string) string {
	return genericType
}

func ColumnsSql(table *Table) string {
	colSql := make([]string, 0, len(table.Columns))

	for _, col := range table.Columns {
		colSql = append(colSql, fmt.Sprintf("%v %v", col.Name, PostgresType(col.Type)))
	}

	return strings.Join(colSql, ",\n\t")
}
<|endoftext|>