Support AI.INFO (#7)

Author: dengliming

redisai/commands.go

@@ -1,7 +1,10 @@
 package redisai
 
 import (
+	"errors"
+	"fmt"
 	"github.com/gomodule/redigo/redis"
+	"strconv"
 )
 
 // TensorSet sets a tensor
@@ -163,3 +166,35 @@ func (c *Client) LoadBackend(backend_identifier string, location string) (err er
 	_, err = c.DoOrSend("AI.CONFIG", args, nil)
 	return
 }
+
+// Returns information about the execution a model or a script.
+func (c *Client) Info(key string) (map[string]string, error) {
+	reply, err := c.DoOrSend("AI.INFO", redis.Args{key}, nil)
+	values, err := redis.Values(reply, err)
+	if err != nil {
+		return nil, err
+	}
+	if len(values)%2 != 0 {
+		return nil, errors.New("expects even number of values result")
+	}
+
+	m := make(map[string]string, len(values)/2)
+	for i := 0; i < len(values); i += 2 {
+		k := string(values[i].([]byte))
+		switch v := values[i+1].(type) {
+		case []byte:
+			m[k] = string(values[i+1].([]byte))
+			break
+		case int64:
+			m[k] = strconv.FormatInt(values[i+1].(int64), 10)
+		default:
+			return nil, fmt.Errorf("unexpected element type for (Ints,String), got type %T", v)
+		}
+	}
+	return m, nil
+}
+
+// Resets all statistics associated with the key
+func (c *Client) ResetStat(key string) (string, error) {
+	return redis.String(c.DoOrSend("AI.INFO", redis.Args{key, "RESETSTAT"}, nil))
+}

redisai/commands_test.go

@@ -825,3 +825,46 @@ func TestCommand_LoadBackend(t *testing.T) {
 		})
 	}
 }
+
+func TestCommand_Info(t *testing.T) {
+	c := createTestClient()
+	keyModel1 := "test:Info:1"
+	data, err := ioutil.ReadFile("./../tests/test_data/graph.pb")
+	if err != nil {
+		t.Errorf("Error preparing for Info(), while issuing ModelSet. error = %v", err)
+		return
+	}
+	err = c.ModelSet(keyModel1, BackendTF, DeviceCPU, data, []string{"a", "b"}, []string{"mul"})
+	if err != nil {
+		t.Errorf("Error preparing for Info(), while issuing ModelSet. error = %v", err)
+		return
+	}
+
+	// first inited info
+	info, err := c.Info(keyModel1)
+	assert.NotNil(t, info)
+	assert.Equal(t, keyModel1, info["key"])
+	assert.Equal(t, DeviceCPU, info["device"])
+	assert.Equal(t, BackendTF, info["backend"])
+	assert.Equal(t, "0", info["calls"])
+
+	err = c.TensorSet("a", TypeFloat32, []int{1}, []float32{1.1})
+	assert.Nil(t, err)
+	err = c.TensorSet("b", TypeFloat32, []int{1}, []float32{4.4})
+	assert.Nil(t, err)
+	err = c.ModelRun(keyModel1, []string{"a", "b"}, []string{"mul"})
+	assert.Nil(t, err)
+	info, err = c.Info(keyModel1)
+	// one model runs
+	assert.Equal(t, "1", info["calls"])
+
+	// reset
+	ret, err := c.ResetStat(keyModel1)
+	assert.Equal(t, "OK", ret)
+	info, err = c.Info(keyModel1)
+	assert.Equal(t, "0", info["calls"])
+
+	// not exits
+	ret, err = c.ResetStat("notExits")
+	assert.NotNil(t, err)
+}