framework.go

package dlframework

import (
  "fmt"
  "io"
  "io/ioutil"
  "net/http"
  "os"
  "path/filepath"
  "strings"

  "github.com/Masterminds/semver"
  "github.com/pkg/errors"
  "github.com/unknwon/com"
  "golang.org/x/sync/syncmap"

  "github.com/c3sr/config"
  yaml "gopkg.in/yaml.v2"
)

var frameworkRegistry = syncmap.Map{}

func (f FrameworkManifest) MustCanonicalName() string {
  s, err := f.CanonicalName()
  if err != nil {
    log.WithField("framework_name", f.GetName()).Fatal("unable to get framework canonical name")
    return ""
  }
  return s
}

func (f FrameworkManifest) CanonicalName() (string, error) {
  frameworkName := f.GetName()
  if frameworkName == "" {
    return "", errors.New("framework name must not be empty")
  }
  frameworkVersion := f.GetVersion()
  if frameworkVersion == "" {
    frameworkVersion = "latest"
  }
  return CleanString(frameworkName) + ":" + CleanString(frameworkVersion), nil
}

func (f FrameworkManifest) Register() error {
  name, err := f.CanonicalName()
  if err != nil {
    return err
  }
  return f.RegisterNamed(name)
}

func (f FrameworkManifest) RegisterNamed(s string) error {
  name := CleanString(s)
  if _, ok := frameworkRegistry.Load(name); ok {
    return errors.Errorf("the %s framework has already been registered", s)
  }
  frameworkRegistry.Store(s, f)
  return nil
}

func RegisteredFrameworkNames() []string {
  return syncMapKeys(&frameworkRegistry)
}

func Frameworks() ([]FrameworkManifest, error) {
  names := RegisteredFrameworkNames()
  fws := make([]FrameworkManifest, len(names))
  for ii, name := range names {
    f, ok := frameworkRegistry.Load(name)
    if !ok {
      return nil, errors.Errorf("framework %s was not found", name)
    }
    fw, ok := f.(FrameworkManifest)
    if !ok {
      return nil, errors.Errorf("framework %s was found but not of type FrameworkManifest", name)
    }
    fws[ii] = fw
  }
  return fws, nil
}

func (f FrameworkManifest) Models() []ModelManifest {
  models := []ModelManifest{}
  wDir, err := f.WorkDir()
  if err != nil {
    return models
  }

  files, err := ioutil.ReadDir(wDir)
  if err != nil {
    return models
  }

  for _, file := range files {
    m, err := f.FindModel(file.Name())
    if err != nil {
      continue
    }
    models = append(models, *m)
  }

  return models
}

func (f FrameworkManifest) FindModel(name string) (*ModelManifest, error) {
  name = CleanString(name)
  name = strings.Replace(name, ":", "_", -1)

  model := new(ModelManifest)
  frameworkVersionString := f.GetVersion()
  if frameworkVersionString == "" {
    return nil, errors.Errorf("expecting a framework version for framework = %v", f.GetName())
  }

  frameworkVersion, err := semver.NewVersion(frameworkVersionString)
  if err != nil {
    return nil, errors.Wrapf(err, "unable to parse version information for framework = %v", f.GetName())
  }

  wDir, err := f.WorkDir()
  if err != nil {
    return nil, errors.Wrapf(err, "unable to access workDir")
  }

  file := filepath.Join(wDir, name, name+".yml")

  _, err = os.Stat(file)
  if err != nil {
    downloadFile(file, f.Name)
  }

  content, err := ioutil.ReadFile(file)
  if err != nil {
    return nil, errors.Wrapf(err, "unable to read model manifest file %v", file)
  }

  if err := yaml.Unmarshal(content, model); err != nil {
    return nil, errors.Wrapf(err, "unable to unmarshal %v", name)
  }

  if model.GetFramework().GetVersion() != "latest" {
    cs, err := model.FrameworkConstraint()
    if err != nil {
      return nil, errors.Wrapf(err, "unable to parse framework version constraint from %v", name)
    }

    ok := cs.Check(frameworkVersion)
    if !ok {
      return nil, errors.Wrapf(err, "constraint fails, framework = %v, model requires = %v", frameworkVersion, cs)
    }
  }

  return model, nil
}

func (f FrameworkManifest) GetModelManifestFromUrl(url string) (*ModelManifest, error) {
  data, err := http.Get(url)
  if err != nil {
    return nil, err
  }
  defer data.Body.Close()

  if data.StatusCode != http.StatusOK {
    return nil, errors.New(fmt.Sprintf("Attempting to download", url, "but got status", data.StatusCode))
  }

  bodyBuffer, err := ioutil.ReadAll(data.Body)
  if err != nil {
    return nil, err
  }

  model := new(ModelManifest)
  if err := yaml.Unmarshal(bodyBuffer, model); err != nil {
    return nil, errors.Wrapf(err, "unable to unmarshal ModelManifest from %s", url)
  }

  return model, nil
}

func downloadFile(fullFilePath string, frameworkName string) error {
  fmt.Println("Could not find", fullFilePath, ", attempting to download it")
  fileDirectory, filename := filepath.Split(fullFilePath)
  urlToGetModelManifest := generateUrlToDownloadManifest(filename, frameworkName)
  fmt.Println("Downloading ", urlToGetModelManifest)

  err := os.MkdirAll(fileDirectory, os.ModePerm)
  if err != nil {
    return err
  }
  output, err := os.Create(fullFilePath)
  if err != nil {
    return err
  }

  data, err := http.Get(urlToGetModelManifest)
  if err != nil {
    return err
  }
  defer data.Body.Close()

  if data.StatusCode != http.StatusOK {
    return errors.New(fmt.Sprintf("Attempting to download", urlToGetModelManifest, "but got status", data.StatusCode))
  }
  _, err = io.Copy(output, data.Body)
  if err != nil {
    fmt.Println(err)
    return err
  }
  return nil
}

func generateUrlToDownloadManifest(filenameWithExtension string, frameworkName string) string {
  baseUrl := os.Getenv("MODELMANIFEST_BASEURL")
  if baseUrl == "" {
    baseUrl = "https://s3.amazonaws.com/store.carml.org/ModelManifest/"
  }
  urlToGetFile := baseUrl + frameworkName + "/" + filenameWithExtension
  return urlToGetFile
}

func FindFramework(name string) (*FrameworkManifest, error) {
  var framework *FrameworkManifest
  name = CleanString(name)
  frameworkRegistry.Range(func(key0 interface{}, value interface{}) bool {
    key, ok := key0.(string)
    if !ok {
      return true
    }
    if key != name {
      return true
    }
    f, ok := value.(FrameworkManifest)
    if !ok {
      return true
    }
    framework = &f
    return false
  })
  if framework == nil {
    return nil, errors.Errorf("framework %s not found in registry", name)
  }
  return framework, nil
}

func (f *FrameworkManifest) WorkDir() (string, error) {
  workDir := filepath.Join(config.App.TempDir, "models", CleanString(f.GetName()))
  if !com.IsDir(workDir) {
    if err := os.MkdirAll(workDir, 0700); err != nil {
      return "", errors.Wrapf(err, "failed to create framework work directory %v", workDir)
    }
  }
  return workDir, nil
}