Container request priority flag.

[lightning.git] / import.go

diff --git a/import.go b/import.go
index 49a7c7a885b2e1478a6b31b9e5bcda4cdbdd24fd..e6b86afd81ca22038c025c669e39a6fd3a0c3b88 100644 (file)
--- a/import.go
+++ b/import.go
@@ -8,7 +8,6 @@ import (
        "flag"
        "fmt"
        "io"
-       "log"
        "net/http"
        _ "net/http/pprof"
        "os"
@@ -21,12 +20,17 @@ import (
        "sync"
        "sync/atomic"
        "time"
+
+       "git.arvados.org/arvados.git/sdk/go/arvados"
+       log "github.com/sirupsen/logrus"
 )
 
 type importer struct {
        tagLibraryFile string
        refFile        string
        outputFile     string
+       projectUUID    string
+       runLocal       bool
        encoder        *gob.Encoder
 }
 
@@ -41,7 +45,10 @@ func (cmd *importer) RunCommand(prog string, args []string, stdin io.Reader, std
        flags.SetOutput(stderr)
        flags.StringVar(&cmd.tagLibraryFile, "tag-library", "", "tag library fasta `file`")
        flags.StringVar(&cmd.refFile, "ref", "", "reference fasta `file`")
-       flags.StringVar(&cmd.outputFile, "o", "", "output `file`")
+       flags.StringVar(&cmd.outputFile, "o", "-", "output `file`")
+       flags.StringVar(&cmd.projectUUID, "project", "", "project `UUID` for output data")
+       flags.BoolVar(&cmd.runLocal, "local", false, "run on local host (default: run in an arvados container)")
+       priority := flags.Int("priority", 500, "container request priority")
        pprof := flags.String("pprof", "", "serve Go profile data at http://`[addr]:port`")
        err = flags.Parse(args)
        if err == flag.ErrHelp {
@@ -63,6 +70,45 @@ func (cmd *importer) RunCommand(prog string, args []string, stdin io.Reader, std
                }()
        }
 
+       if !cmd.runLocal {
+               runner := arvadosContainerRunner{
+                       Name:        "lightning import",
+                       Client:      arvados.NewClientFromEnv(),
+                       ProjectUUID: cmd.projectUUID,
+                       RAM:         30000000000,
+                       VCPUs:       16,
+                       Priority:    *priority,
+               }
+               err = runner.TranslatePaths(&cmd.tagLibraryFile, &cmd.refFile, &cmd.outputFile)
+               if err != nil {
+                       return 1
+               }
+               inputs := flags.Args()
+               for i := range inputs {
+                       err = runner.TranslatePaths(&inputs[i])
+                       if err != nil {
+                               return 1
+                       }
+               }
+               if cmd.outputFile == "-" {
+                       cmd.outputFile = "/mnt/output/library.gob"
+               } else {
+                       // Not yet implemented, but this should write
+                       // the collection to an existing collection,
+                       // possibly even an in-place update.
+                       err = errors.New("cannot specify output file in container mode: not implemented")
+                       return 1
+               }
+               runner.Args = append([]string{"import", "-local=true", "-tag-library", cmd.tagLibraryFile, "-ref", cmd.refFile, "-o", cmd.outputFile}, inputs...)
+               var output string
+               output, err = runner.Run()
+               if err != nil {
+                       return 1
+               }
+               fmt.Fprintln(stdout, output+"/library.gob")
+               return 0
+       }
+
        infiles, err := listInputFiles(flags.Args())
        if err != nil {
                return 1
@@ -73,38 +119,35 @@ func (cmd *importer) RunCommand(prog string, args []string, stdin io.Reader, std
                return 1
        }
        go func() {
-               for range time.Tick(10 * time.Second) {
+               for range time.Tick(10 * time.Minute) {
                        log.Printf("tilelib.Len() == %d", tilelib.Len())
                }
        }()
 
-       var outfile *os.File
-       var w *bufio.Writer
-       if cmd.outputFile == "" {
-               w = bufio.NewWriter(stdout)
+       var output io.WriteCloser
+       if cmd.outputFile == "-" {
+               output = nopCloser{stdout}
        } else {
-               outfile, err = os.OpenFile(cmd.outputFile, os.O_CREATE|os.O_WRONLY, 0777)
+               output, err = os.OpenFile(cmd.outputFile, os.O_CREATE|os.O_WRONLY, 0777)
                if err != nil {
                        return 1
                }
-               defer outfile.Close()
-               w = bufio.NewWriter(outfile)
+               defer output.Close()
        }
-       cmd.encoder = gob.NewEncoder(w)
+       bufw := bufio.NewWriter(output)
+       cmd.encoder = gob.NewEncoder(bufw)
 
        err = cmd.tileInputs(tilelib, infiles)
        if err != nil {
                return 1
        }
-       err = w.Flush()
+       err = bufw.Flush()
        if err != nil {
                return 1
        }
-       if outfile != nil {
-               err = outfile.Close()
-               if err != nil {
-                       return 1
-               }
+       err = output.Close()
+       if err != nil {
+               return 1
        }
        return 0
 }
@@ -266,9 +309,10 @@ func (cmd *importer) tileInputs(tilelib *tileLibrary, infiles []string) error {
        }
        go close(todo)
        var tileJobs sync.WaitGroup
-       running := int64(runtime.NumCPU())
-       for i := 0; i < runtime.NumCPU(); i++ {
+       var running int64
+       for i := 0; i < runtime.NumCPU()*9/8+1; i++ {
                tileJobs.Add(1)
+               atomic.AddInt64(&running, 1)
                go func() {
                        defer tileJobs.Done()
                        defer atomic.AddInt64(&running, -1)