12630: Call nvidia-modprobe, support CUDA_VISIBLE_DEVICES
[arvados.git] / lib / crunchrun / singularity.go
1 // Copyright (C) The Arvados Authors. All rights reserved.
2 //
3 // SPDX-License-Identifier: AGPL-3.0
4
5 package crunchrun
6
7 import (
8         "fmt"
9         "io/ioutil"
10         "os"
11         "os/exec"
12         "sort"
13         "strings"
14         "syscall"
15         "time"
16
17         "git.arvados.org/arvados.git/sdk/go/arvados"
18         "golang.org/x/net/context"
19 )
20
21 type singularityExecutor struct {
22         logf          func(string, ...interface{})
23         spec          containerSpec
24         tmpdir        string
25         child         *exec.Cmd
26         imageFilename string // "sif" image
27 }
28
29 func newSingularityExecutor(logf func(string, ...interface{})) (*singularityExecutor, error) {
30         tmpdir, err := ioutil.TempDir("", "crunch-run-singularity-")
31         if err != nil {
32                 return nil, err
33         }
34         return &singularityExecutor{
35                 logf:   logf,
36                 tmpdir: tmpdir,
37         }, nil
38 }
39
40 func (e *singularityExecutor) Runtime() string { return "singularity" }
41
42 func (e *singularityExecutor) getOrCreateProject(ownerUuid string, name string, containerClient *arvados.Client) (*arvados.Group, error) {
43         var gp arvados.GroupList
44         err := containerClient.RequestAndDecode(&gp,
45                 arvados.EndpointGroupList.Method,
46                 arvados.EndpointGroupList.Path,
47                 nil, arvados.ListOptions{Filters: []arvados.Filter{
48                         arvados.Filter{"owner_uuid", "=", ownerUuid},
49                         arvados.Filter{"name", "=", name},
50                         arvados.Filter{"group_class", "=", "project"},
51                 },
52                         Limit: 1})
53         if err != nil {
54                 return nil, err
55         }
56         if len(gp.Items) == 1 {
57                 return &gp.Items[0], nil
58         }
59
60         var rgroup arvados.Group
61         err = containerClient.RequestAndDecode(&rgroup,
62                 arvados.EndpointGroupCreate.Method,
63                 arvados.EndpointGroupCreate.Path,
64                 nil, map[string]interface{}{
65                         "group": map[string]string{
66                                 "owner_uuid":  ownerUuid,
67                                 "name":        name,
68                                 "group_class": "project",
69                         },
70                 })
71         if err != nil {
72                 return nil, err
73         }
74         return &rgroup, nil
75 }
76
77 func (e *singularityExecutor) checkImageCache(dockerImageID string, container arvados.Container, arvMountPoint string,
78         containerClient *arvados.Client) (collection *arvados.Collection, err error) {
79
80         // Cache the image to keep
81         cacheGroup, err := e.getOrCreateProject(container.RuntimeUserUUID, ".cache", containerClient)
82         if err != nil {
83                 return nil, fmt.Errorf("error getting '.cache' project: %v", err)
84         }
85         imageGroup, err := e.getOrCreateProject(cacheGroup.UUID, "auto-generated singularity images", containerClient)
86         if err != nil {
87                 return nil, fmt.Errorf("error getting 'auto-generated singularity images' project: %s", err)
88         }
89
90         collectionName := fmt.Sprintf("singularity image for %v", dockerImageID)
91         var cl arvados.CollectionList
92         err = containerClient.RequestAndDecode(&cl,
93                 arvados.EndpointCollectionList.Method,
94                 arvados.EndpointCollectionList.Path,
95                 nil, arvados.ListOptions{Filters: []arvados.Filter{
96                         arvados.Filter{"owner_uuid", "=", imageGroup.UUID},
97                         arvados.Filter{"name", "=", collectionName},
98                 },
99                         Limit: 1})
100         if err != nil {
101                 return nil, fmt.Errorf("error querying for collection '%v': %v", collectionName, err)
102         }
103         var imageCollection arvados.Collection
104         if len(cl.Items) == 1 {
105                 imageCollection = cl.Items[0]
106         } else {
107                 collectionName := "converting " + collectionName
108                 exp := time.Now().Add(24 * 7 * 2 * time.Hour)
109                 err = containerClient.RequestAndDecode(&imageCollection,
110                         arvados.EndpointCollectionCreate.Method,
111                         arvados.EndpointCollectionCreate.Path,
112                         nil, map[string]interface{}{
113                                 "collection": map[string]string{
114                                         "owner_uuid": imageGroup.UUID,
115                                         "name":       collectionName,
116                                         "trash_at":   exp.UTC().Format(time.RFC3339),
117                                 },
118                                 "ensure_unique_name": true,
119                         })
120                 if err != nil {
121                         return nil, fmt.Errorf("error creating '%v' collection: %s", collectionName, err)
122                 }
123
124         }
125
126         return &imageCollection, nil
127 }
128
129 // LoadImage will satisfy ContainerExecuter interface transforming
130 // containerImage into a sif file for later use.
131 func (e *singularityExecutor) LoadImage(dockerImageID string, imageTarballPath string, container arvados.Container, arvMountPoint string,
132         containerClient *arvados.Client) error {
133
134         var imageFilename string
135         var sifCollection *arvados.Collection
136         var err error
137         if containerClient != nil {
138                 sifCollection, err = e.checkImageCache(dockerImageID, container, arvMountPoint, containerClient)
139                 if err != nil {
140                         return err
141                 }
142                 imageFilename = fmt.Sprintf("%s/by_uuid/%s/image.sif", arvMountPoint, sifCollection.UUID)
143         } else {
144                 imageFilename = e.tmpdir + "/image.sif"
145         }
146
147         if _, err := os.Stat(imageFilename); os.IsNotExist(err) {
148                 // Make sure the docker image is readable, and error
149                 // out if not.
150                 if _, err := os.Stat(imageTarballPath); err != nil {
151                         return err
152                 }
153
154                 e.logf("building singularity image")
155                 // "singularity build" does not accept a
156                 // docker-archive://... filename containing a ":" character,
157                 // as in "/path/to/sha256:abcd...1234.tar". Workaround: make a
158                 // symlink that doesn't have ":" chars.
159                 err := os.Symlink(imageTarballPath, e.tmpdir+"/image.tar")
160                 if err != nil {
161                         return err
162                 }
163
164                 // Set up a cache and tmp dir for singularity build
165                 err = os.Mkdir(e.tmpdir+"/cache", 0700)
166                 if err != nil {
167                         return err
168                 }
169                 defer os.RemoveAll(e.tmpdir + "/cache")
170                 err = os.Mkdir(e.tmpdir+"/tmp", 0700)
171                 if err != nil {
172                         return err
173                 }
174                 defer os.RemoveAll(e.tmpdir + "/tmp")
175
176                 build := exec.Command("singularity", "build", imageFilename, "docker-archive://"+e.tmpdir+"/image.tar")
177                 build.Env = os.Environ()
178                 build.Env = append(build.Env, "SINGULARITY_CACHEDIR="+e.tmpdir+"/cache")
179                 build.Env = append(build.Env, "SINGULARITY_TMPDIR="+e.tmpdir+"/tmp")
180                 e.logf("%v", build.Args)
181                 out, err := build.CombinedOutput()
182                 // INFO:    Starting build...
183                 // Getting image source signatures
184                 // Copying blob ab15617702de done
185                 // Copying config 651e02b8a2 done
186                 // Writing manifest to image destination
187                 // Storing signatures
188                 // 2021/04/22 14:42:14  info unpack layer: sha256:21cbfd3a344c52b197b9fa36091e66d9cbe52232703ff78d44734f85abb7ccd3
189                 // INFO:    Creating SIF file...
190                 // INFO:    Build complete: arvados-jobs.latest.sif
191                 e.logf("%s", out)
192                 if err != nil {
193                         return err
194                 }
195         }
196
197         if containerClient == nil {
198                 e.imageFilename = imageFilename
199                 return nil
200         }
201
202         // update TTL to now + two weeks
203         exp := time.Now().Add(24 * 7 * 2 * time.Hour)
204
205         uuidPath, err := containerClient.PathForUUID("update", sifCollection.UUID)
206         if err != nil {
207                 e.logf("error PathForUUID: %v", err)
208                 return nil
209         }
210         var imageCollection arvados.Collection
211         err = containerClient.RequestAndDecode(&imageCollection,
212                 arvados.EndpointCollectionUpdate.Method,
213                 uuidPath,
214                 nil, map[string]interface{}{
215                         "collection": map[string]string{
216                                 "name":     fmt.Sprintf("singularity image for %v", dockerImageID),
217                                 "trash_at": exp.UTC().Format(time.RFC3339),
218                         },
219                 })
220         if err == nil {
221                 // If we just wrote the image to the cache, the
222                 // response also returns the updated PDH
223                 e.imageFilename = fmt.Sprintf("%s/by_id/%s/image.sif", arvMountPoint, imageCollection.PortableDataHash)
224                 return nil
225         }
226
227         e.logf("error updating/renaming collection for cached sif image: %v", err)
228         // Failed to update but maybe it lost a race and there is
229         // another cached collection in the same place, so check the cache
230         // again
231         sifCollection, err = e.checkImageCache(dockerImageID, container, arvMountPoint, containerClient)
232         if err != nil {
233                 return err
234         }
235         e.imageFilename = fmt.Sprintf("%s/by_id/%s/image.sif", arvMountPoint, sifCollection.PortableDataHash)
236
237         return nil
238 }
239
240 func (e *singularityExecutor) Create(spec containerSpec) error {
241         e.spec = spec
242         return nil
243 }
244
245 func (e *singularityExecutor) execCmd(path string) *exec.Cmd {
246         args := []string{path, "exec", "--containall", "--cleanenv", "--pwd", e.spec.WorkingDir}
247         if !e.spec.EnableNetwork {
248                 args = append(args, "--net", "--network=none")
249         }
250
251         if e.spec.CUDADeviceCount != 0 {
252                 args = append(args, "--nv")
253         }
254
255         readonlyflag := map[bool]string{
256                 false: "rw",
257                 true:  "ro",
258         }
259         var binds []string
260         for path, _ := range e.spec.BindMounts {
261                 binds = append(binds, path)
262         }
263         sort.Strings(binds)
264         for _, path := range binds {
265                 mount := e.spec.BindMounts[path]
266                 if path == e.spec.Env["HOME"] {
267                         // Singularity treates $HOME as special case
268                         args = append(args, "--home", mount.HostPath+":"+path)
269                 } else {
270                         args = append(args, "--bind", mount.HostPath+":"+path+":"+readonlyflag[mount.ReadOnly])
271                 }
272         }
273
274         // This is for singularity 3.5.2. There are some behaviors
275         // that will change in singularity 3.6, please see:
276         // https://sylabs.io/guides/3.7/user-guide/environment_and_metadata.html
277         // https://sylabs.io/guides/3.5/user-guide/environment_and_metadata.html
278         env := make([]string, 0, len(e.spec.Env))
279         for k, v := range e.spec.Env {
280                 if k == "HOME" {
281                         // Singularity treates $HOME as special case, this is handled
282                         // with --home above
283                         continue
284                 }
285                 env = append(env, "SINGULARITYENV_"+k+"="+v)
286         }
287
288         // Singularity always makes all nvidia devices visible to the
289         // container.  If a resource manager such as slurm or LSF told
290         // us to select specific devices we need to propagate that.
291         for _, s := range os.Environ() {
292                 if strings.HasPrefix(s, "CUDA_VISIBLE_DEVICES=") {
293                         env = append(env, "SINGULARITYENV_"+s)
294                 }
295         }
296
297         args = append(args, e.imageFilename)
298         args = append(args, e.spec.Command...)
299
300         return &exec.Cmd{
301                 Path:   path,
302                 Args:   args,
303                 Env:    env,
304                 Stdin:  e.spec.Stdin,
305                 Stdout: e.spec.Stdout,
306                 Stderr: e.spec.Stderr,
307         }
308 }
309
310 func (e *singularityExecutor) Start() error {
311         path, err := exec.LookPath("singularity")
312         if err != nil {
313                 return err
314         }
315         child := e.execCmd(path)
316         err = child.Start()
317         if err != nil {
318                 return err
319         }
320         e.child = child
321         return nil
322 }
323
324 func (e *singularityExecutor) CgroupID() string {
325         return ""
326 }
327
328 func (e *singularityExecutor) Stop() error {
329         if err := e.child.Process.Signal(syscall.Signal(0)); err != nil {
330                 // process already exited
331                 return nil
332         }
333         return e.child.Process.Signal(syscall.SIGKILL)
334 }
335
336 func (e *singularityExecutor) Wait(context.Context) (int, error) {
337         err := e.child.Wait()
338         if err, ok := err.(*exec.ExitError); ok {
339                 return err.ProcessState.ExitCode(), nil
340         }
341         if err != nil {
342                 return 0, err
343         }
344         return e.child.ProcessState.ExitCode(), nil
345 }
346
347 func (e *singularityExecutor) Close() {
348         err := os.RemoveAll(e.tmpdir)
349         if err != nil {
350                 e.logf("error removing temp dir: %s", err)
351         }
352 }