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 }