github.com/sentienttechnologies/studio-go-runner@v0.0.0-20201118202441-6d21f2ced8ee/internal/runner/request.go (about)

     1  // Copyright 2018-2020 (c) Cognizant Digital Business, Evolutionary AI. All rights reserved. Issued under the Apache 2.0 License.
     2  
     3  package runner
     4  
     5  // This file contains the implementation of a message parser for requests
     6  // arriving from studioml queues formatted using JSON.
     7  //
     8  // To parse and unparse this JSON data use the following ...
     9  //
    10  //    r, err := UnmarshalRequest(bytes)
    11  //    bytes, err = r.Marshal()
    12  
    13  import (
    14  	"bytes"
    15  	"encoding/gob"
    16  	"encoding/json"
    17  
    18  	"github.com/dustin/go-humanize"
    19  
    20  	"github.com/go-stack/stack"
    21  	"github.com/jjeffery/kv" // MIT License
    22  )
    23  
    24  // Resource describes the needed resources for a runner task in a data structure that can be
    25  // marshalled as json
    26  //
    27  type Resource struct {
    28  	Cpus   uint   `json:"cpus"`
    29  	Gpus   uint   `json:"gpus"`
    30  	Hdd    string `json:"hdd"`
    31  	Ram    string `json:"ram"`
    32  	GpuMem string `json:"gpuMem"`
    33  }
    34  
    35  func (rsc Resource) String() (serialized string) {
    36  	serialize, _ := json.Marshal(rsc)
    37  
    38  	return string(serialize)
    39  }
    40  
    41  // Fit determines is a supplied resource description acting as a request can
    42  // be satisfied by the receiver resource
    43  //
    44  func (rsc *Resource) Fit(r *Resource) (didFit bool, err kv.Error) {
    45  
    46  	lRam, errGo := humanize.ParseBytes(rsc.Ram)
    47  	if errGo != nil {
    48  		return false, kv.NewError("left side RAM could not be parsed").With("stack", stack.Trace().TrimRuntime())
    49  	}
    50  
    51  	rRam, errGo := humanize.ParseBytes(r.Ram)
    52  	if errGo != nil {
    53  		return false, kv.NewError("right side RAM could not be parsed").With("stack", stack.Trace().TrimRuntime())
    54  	}
    55  
    56  	lHdd, errGo := humanize.ParseBytes(rsc.Hdd)
    57  	if errGo != nil {
    58  		return false, kv.NewError("left side Hdd could not be parsed").With("stack", stack.Trace().TrimRuntime())
    59  	}
    60  
    61  	rHdd, errGo := humanize.ParseBytes(r.Hdd)
    62  	if errGo != nil {
    63  		return false, kv.NewError("right side Hdd could not be parsed").With("stack", stack.Trace().TrimRuntime())
    64  	}
    65  
    66  	lGpuMem, errGo := humanize.ParseBytes(rsc.GpuMem)
    67  	// GpuMem is optional so handle the case when it does not parse and is empty
    68  	if 0 != len(rsc.GpuMem) {
    69  		if errGo != nil {
    70  			return false, kv.NewError("left side gpuMem could not be parsed").With("left_mem", rsc.GpuMem).With("stack", stack.Trace().TrimRuntime())
    71  		}
    72  	}
    73  
    74  	rGpuMem, errGo := humanize.ParseBytes(r.GpuMem)
    75  	// GpuMem is optional so handle the case when it does not parse and is empty
    76  	if 0 != len(r.GpuMem) {
    77  		if errGo != nil {
    78  			return false, kv.NewError("right side gpuMem could not be parsed").With("right", r.GpuMem).With("stack", stack.Trace().TrimRuntime())
    79  		}
    80  	}
    81  
    82  	return rsc.Cpus <= r.Cpus && rsc.Gpus <= r.Gpus && lHdd <= rHdd && lRam <= rRam && lGpuMem <= rGpuMem, nil
    83  }
    84  
    85  // Clone will deep copy a resource and return the copy
    86  //
    87  func (rsc *Resource) Clone() (r *Resource) {
    88  
    89  	mod := bytes.Buffer{}
    90  	enc := gob.NewEncoder(&mod)
    91  	dec := gob.NewDecoder(&mod)
    92  
    93  	if err := enc.Encode(rsc); err != nil {
    94  		return nil
    95  	}
    96  
    97  	r = &Resource{}
    98  	if err := dec.Decode(r); err != nil {
    99  		return nil
   100  	}
   101  	return r
   102  }
   103  
   104  // Config is a marshalled data structure used with studioml requests for defining the
   105  // configuration of an environment used to run jobs
   106  type Config struct {
   107  	Cloud                  interface{}       `json:"cloud"`
   108  	Database               Database          `json:"database"`
   109  	SaveWorkspaceFrequency string            `json:"saveWorkspaceFrequency"`
   110  	Lifetime               string            `json:"experimentLifetime"`
   111  	Verbose                string            `json:"verbose"`
   112  	Env                    map[string]string `json:"env"`
   113  	Pip                    []string          `json:"pip"`
   114  	Runner                 RunnerCustom      `json:"runner"`
   115  }
   116  
   117  // RunnerCustom defines a custom type of resource used by the go runner to implement a slack
   118  // notification mechanism
   119  //
   120  type RunnerCustom struct {
   121  	SlackDest string `json:"slack_destination"`
   122  }
   123  
   124  // Database marshalls the studioML database specification for experiment meta data
   125  type Database struct {
   126  	ApiKey            string `json:"apiKey"`
   127  	AuthDomain        string `json:"authDomain"`
   128  	DatabaseURL       string `json:"databaseURL"`
   129  	MessagingSenderId int64  `json:"messagingSenderId"`
   130  	ProjectId         string `json:"projectId"`
   131  	StorageBucket     string `json:"storageBucket"`
   132  	Type              string `json:"type"`
   133  	UseEmailAuth      bool   `json:"use_email_auth"`
   134  }
   135  
   136  // Experiment marshalls the studioML experiment meta data
   137  type Experiment struct {
   138  	Args               []string            `json:"args"`
   139  	Artifacts          map[string]Artifact `json:"artifacts"`
   140  	Filename           string              `json:"filename"`
   141  	Git                interface{}         `json:"git"`
   142  	Info               Info                `json:"info"`
   143  	Key                string              `json:"key"`
   144  	Metric             interface{}         `json:"metric"`
   145  	Project            interface{}         `json:"project"`
   146  	Pythonenv          []string            `json:"pythonenv"`
   147  	PythonVer          string              `json:"pythonver"`
   148  	Resource           Resource            `json:"resources_needed"`
   149  	Status             string              `json:"status"`
   150  	TimeAdded          float64             `json:"time_added"`
   151  	MaxDuration        string              `json:"max_duration"`
   152  	TimeFinished       interface{}         `json:"time_finished"`
   153  	TimeLastCheckpoint interface{}         `json:"time_last_checkpoint"`
   154  	TimeStarted        interface{}         `json:"time_started"`
   155  }
   156  
   157  // Request marshalls the requests made by studioML under which all of the other
   158  // meta data can be found
   159  type Request struct {
   160  	Config     Config     `json:"config"`
   161  	Experiment Experiment `json:"experiment"`
   162  }
   163  
   164  // Info is a marshalled item from the studioML experiment definition that
   165  // is ignored by the go runner and so is stubbed out
   166  type Info struct {
   167  }
   168  
   169  // Artifact is a marshalled component of a StudioML experiment definition that
   170  // is used to encapsulate files and other external data sources
   171  // that the runner retrieve and/or upload as the experiment progresses
   172  type Artifact struct {
   173  	Bucket    string `json:"bucket"`
   174  	Key       string `json:"key"`
   175  	Hash      string `json:"hash,omitempty"`
   176  	Local     string `json:"local,omitempty"`
   177  	Mutable   bool   `json:"mutable"`
   178  	Unpack    bool   `json:"unpack"`
   179  	Qualified string `json:"qualified"`
   180  }
   181  
   182  // Clone is a full on duplication of the original artifact
   183  func (a *Artifact) Clone() (b *Artifact) {
   184  	return &Artifact{
   185  		Bucket:    a.Bucket[:],
   186  		Key:       a.Key[:],
   187  		Hash:      a.Hash[:],
   188  		Local:     a.Local[:],
   189  		Mutable:   a.Mutable,
   190  		Unpack:    a.Unpack,
   191  		Qualified: a.Qualified[:],
   192  	}
   193  }
   194  
   195  // UnmarshalRequest takes an encoded StudioML request and extracts it
   196  // into go data structures used by the go runner
   197  //
   198  func UnmarshalRequest(data []byte) (r *Request, err kv.Error) {
   199  	r = &Request{}
   200  	errGo := json.Unmarshal(data, r)
   201  	if errGo != nil {
   202  		return nil, kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime())
   203  	}
   204  	return r, nil
   205  }
   206  
   207  // Marshal takes the go data structure used to define a StudioML experiment
   208  // request and serializes it as json to the byte array
   209  //
   210  func (r *Request) Marshal() (buffer []byte, err kv.Error) {
   211  	buffer, errGo := json.Marshal(r)
   212  	if errGo != nil {
   213  		return nil, kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime())
   214  	}
   215  	return buffer, nil
   216  }