Fix some tests.
[lightning.git] / import.go
index 0d2d7379d184174a54766a9a4b4306adb14fb582..20058dc8a9f184dbc5c88c9d32cd32399dd3da34 100644 (file)
--- a/import.go
+++ b/import.go
@@ -1,14 +1,20 @@
-package main
+// Copyright (C) The Lightning Authors. All rights reserved.
+//
+// SPDX-License-Identifier: AGPL-3.0
+
+package lightning
 
 import (
        "bufio"
        "compress/gzip"
+       "context"
        "encoding/gob"
        "encoding/json"
        "errors"
        "flag"
        "fmt"
        "io"
+       "io/ioutil"
        "net/http"
        _ "net/http/pprof"
        "os"
@@ -22,13 +28,8 @@ import (
        "sync/atomic"
        "time"
 
-       "git.arvados.org/arvados.git/sdk/go/arvados"
-       "github.com/lucasb-eyer/go-colorful"
+       "github.com/klauspost/pgzip"
        log "github.com/sirupsen/logrus"
-       "gonum.org/v1/plot"
-       "gonum.org/v1/plot/plotter"
-       "gonum.org/v1/plot/vg"
-       "gonum.org/v1/plot/vg/draw"
 )
 
 type importer struct {
@@ -36,6 +37,8 @@ type importer struct {
        refFile             string
        outputFile          string
        projectUUID         string
+       loglevel            string
+       priority            int
        runLocal            bool
        skipOOO             bool
        outputTiles         bool
@@ -43,6 +46,8 @@ type importer struct {
        outputStats         string
        matchChromosome     *regexp.Regexp
        encoder             *gob.Encoder
+       retainAfterEncoding bool // keep imported genomes/refseqs in memory after writing to disk
+       batchArgs
 }
 
 func (cmd *importer) RunCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
@@ -63,10 +68,11 @@ func (cmd *importer) RunCommand(prog string, args []string, stdin io.Reader, std
        flags.BoolVar(&cmd.outputTiles, "output-tiles", false, "include tile variant sequences in output file")
        flags.BoolVar(&cmd.saveIncompleteTiles, "save-incomplete-tiles", false, "treat tiles with no-calls as regular tiles")
        flags.StringVar(&cmd.outputStats, "output-stats", "", "output stats to `file` (json)")
+       cmd.batchArgs.Flags(flags)
        matchChromosome := flags.String("match-chromosome", "^(chr)?([0-9]+|X|Y|MT?)$", "import chromosomes that match the given `regexp`")
-       priority := flags.Int("priority", 500, "container request priority")
+       flags.IntVar(&cmd.priority, "priority", 500, "container request priority")
        pprof := flags.String("pprof", "", "serve Go profile data at http://`[addr]:port`")
-       loglevel := flags.String("loglevel", "info", "logging threshold (trace, debug, info, warn, error, fatal, or panic)")
+       flags.StringVar(&cmd.loglevel, "loglevel", "info", "logging threshold (trace, debug, info, warn, error, fatal, or panic)")
        err = flags.Parse(args)
        if err == flag.ErrHelp {
                err = nil
@@ -87,7 +93,7 @@ func (cmd *importer) RunCommand(prog string, args []string, stdin io.Reader, std
                }()
        }
 
-       lvl, err := log.ParseLevel(*loglevel)
+       lvl, err := log.ParseLevel(cmd.loglevel)
        if err != nil {
                return 2
        }
@@ -99,52 +105,10 @@ 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:         80000000000,
-                       VCPUs:       32,
-                       Priority:    *priority,
-               }
-               err = runner.TranslatePaths(&cmd.tagLibraryFile, &cmd.refFile, &cmd.outputFile)
+               err = cmd.runBatches(stdout, flags.Args())
                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.gz"
-               } 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",
-                       "-loglevel=" + *loglevel,
-                       fmt.Sprintf("-skip-ooo=%v", cmd.skipOOO),
-                       fmt.Sprintf("-output-tiles=%v", cmd.outputTiles),
-                       fmt.Sprintf("-save-incomplete-tiles=%v", cmd.saveIncompleteTiles),
-                       "-match-chromosome", cmd.matchChromosome.String(),
-                       "-output-stats", "/mnt/output/stats.json",
-                       "-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.gz")
                return 0
        }
 
@@ -152,6 +116,7 @@ func (cmd *importer) RunCommand(prog string, args []string, stdin io.Reader, std
        if err != nil {
                return 1
        }
+       infiles = cmd.batchArgs.Slice(infiles)
 
        taglib, err := cmd.loadTagLibrary()
        if err != nil {
@@ -168,12 +133,12 @@ func (cmd *importer) RunCommand(prog string, args []string, stdin io.Reader, std
                }
                defer outf.Close()
                if strings.HasSuffix(cmd.outputFile, ".gz") {
-                       outw = gzip.NewWriter(outf)
+                       outw = pgzip.NewWriter(outf)
                } else {
                        outw = outf
                }
        }
-       bufw := bufio.NewWriter(outw)
+       bufw := bufio.NewWriterSize(outw, 64*1024*1024)
        cmd.encoder = gob.NewEncoder(bufw)
 
        tilelib := &tileLibrary{taglib: taglib, retainNoCalls: cmd.saveIncompleteTiles, skipOOO: cmd.skipOOO}
@@ -208,33 +173,95 @@ func (cmd *importer) RunCommand(prog string, args []string, stdin io.Reader, std
        return 0
 }
 
-func (cmd *importer) tileFasta(tilelib *tileLibrary, infile string) (tileSeq, []importStats, error) {
+func (cmd *importer) runBatches(stdout io.Writer, inputs []string) error {
+       if cmd.outputFile != "-" {
+               // Not yet implemented, but this should write
+               // the collection to an existing collection,
+               // possibly even an in-place update.
+               return errors.New("cannot specify output file in container mode: not implemented")
+       }
+       runner := arvadosContainerRunner{
+               Name:        "lightning import",
+               Client:      arvadosClientFromEnv,
+               ProjectUUID: cmd.projectUUID,
+               APIAccess:   true,
+               RAM:         350000000000,
+               VCPUs:       96,
+               Priority:    cmd.priority,
+               KeepCache:   1,
+       }
+       err := runner.TranslatePaths(&cmd.tagLibraryFile, &cmd.refFile, &cmd.outputFile)
+       if err != nil {
+               return err
+       }
+       for i := range inputs {
+               err = runner.TranslatePaths(&inputs[i])
+               if err != nil {
+                       return err
+               }
+       }
+
+       outputs, err := cmd.batchArgs.RunBatches(context.Background(), func(ctx context.Context, batch int) (string, error) {
+               runner := runner
+               if cmd.batches > 1 {
+                       runner.Name += fmt.Sprintf(" (batch %d of %d)", batch, cmd.batches)
+               }
+               runner.Args = []string{"import",
+                       "-local=true",
+                       "-loglevel=" + cmd.loglevel,
+                       "-pprof=:6061",
+                       fmt.Sprintf("-skip-ooo=%v", cmd.skipOOO),
+                       fmt.Sprintf("-output-tiles=%v", cmd.outputTiles),
+                       fmt.Sprintf("-save-incomplete-tiles=%v", cmd.saveIncompleteTiles),
+                       "-match-chromosome", cmd.matchChromosome.String(),
+                       "-output-stats", "/mnt/output/stats.json",
+                       "-tag-library", cmd.tagLibraryFile,
+                       "-ref", cmd.refFile,
+                       "-o", "/mnt/output/library.gob.gz",
+               }
+               runner.Args = append(runner.Args, cmd.batchArgs.Args(batch)...)
+               runner.Args = append(runner.Args, inputs...)
+               return runner.RunContext(ctx)
+       })
+       if err != nil {
+               return err
+       }
+       var outfiles []string
+       for _, o := range outputs {
+               outfiles = append(outfiles, o+"/library.gob.gz")
+       }
+       fmt.Fprintln(stdout, strings.Join(outfiles, " "))
+       return nil
+}
+
+func (cmd *importer) tileFasta(tilelib *tileLibrary, infile string, isRef bool) (tileSeq, []importStats, error) {
        var input io.ReadCloser
-       input, err := os.Open(infile)
+       input, err := open(infile)
        if err != nil {
                return nil, nil, err
        }
        defer input.Close()
+       input = ioutil.NopCloser(bufio.NewReaderSize(input, 8*1024*1024))
        if strings.HasSuffix(infile, ".gz") {
-               input, err = gzip.NewReader(input)
+               input, err = pgzip.NewReader(input)
                if err != nil {
                        return nil, nil, err
                }
                defer input.Close()
        }
-       return tilelib.TileFasta(infile, input, cmd.matchChromosome)
+       return tilelib.TileFasta(infile, input, cmd.matchChromosome, isRef)
 }
 
 func (cmd *importer) loadTagLibrary() (*tagLibrary, error) {
        log.Printf("tag library %s load starting", cmd.tagLibraryFile)
-       f, err := os.Open(cmd.tagLibraryFile)
+       f, err := open(cmd.tagLibraryFile)
        if err != nil {
                return nil, err
        }
        defer f.Close()
-       var rdr io.ReadCloser = f
+       rdr := ioutil.NopCloser(bufio.NewReaderSize(f, 64*1024*1024))
        if strings.HasSuffix(cmd.tagLibraryFile, ".gz") {
-               rdr, err = gzip.NewReader(f)
+               rdr, err = gzip.NewReader(rdr)
                if err != nil {
                        return nil, fmt.Errorf("%s: gzip: %s", cmd.tagLibraryFile, err)
                }
@@ -254,8 +281,8 @@ func (cmd *importer) loadTagLibrary() (*tagLibrary, error) {
 
 var (
        vcfFilenameRe    = regexp.MustCompile(`\.vcf(\.gz)?$`)
-       fasta1FilenameRe = regexp.MustCompile(`\.1\.fa(sta)?(\.gz)?$`)
-       fasta2FilenameRe = regexp.MustCompile(`\.2\.fa(sta)?(\.gz)?$`)
+       fasta1FilenameRe = regexp.MustCompile(`\.1\.fa(sta)?(\.fa(sta)?)?(\.gz)?$`)
+       fasta2FilenameRe = regexp.MustCompile(`\.2\.fa(sta)?(\.fa(sta)?)?(\.gz)?$`)
        fastaFilenameRe  = regexp.MustCompile(`\.fa(sta)?(\.gz)?$`)
 )
 
@@ -320,35 +347,34 @@ func (cmd *importer) tileInputs(tilelib *tileLibrary, infiles []string) error {
                if fasta1FilenameRe.MatchString(infile) {
                        todo <- func() error {
                                defer phases.Done()
-                               log.Printf("%s starting", infile)
+                               log.Printf("%s (sample.1) starting tiling", infile)
                                defer log.Printf("%s done", infile)
-                               tseqs, stats, err := cmd.tileFasta(tilelib, infile)
+                               tseqs, stats, err := cmd.tileFasta(tilelib, infile, false)
                                allstats[idx*2] = stats
                                var kept, dropped int
                                variants[0], kept, dropped = tseqs.Variants()
-                               log.Printf("%s found %d unique tags plus %d repeats", infile, kept, dropped)
+                               log.Printf("%s (sample.1) found %d unique tags plus %d repeats", infile, kept, dropped)
                                return err
                        }
-                       infile2 := fasta1FilenameRe.ReplaceAllString(infile, `.2.fa$1$2`)
+                       infile2 := fasta1FilenameRe.ReplaceAllString(infile, `.2.fa$1$2$4`)
                        todo <- func() error {
                                defer phases.Done()
-                               log.Printf("%s starting", infile2)
+                               log.Printf("%s (sample.2) starting tiling", infile2)
                                defer log.Printf("%s done", infile2)
-                               tseqs, stats, err := cmd.tileFasta(tilelib, infile2)
+                               tseqs, stats, err := cmd.tileFasta(tilelib, infile2, false)
                                allstats[idx*2+1] = stats
                                var kept, dropped int
                                variants[1], kept, dropped = tseqs.Variants()
-                               log.Printf("%s found %d unique tags plus %d repeats", infile2, kept, dropped)
-
+                               log.Printf("%s (sample.2) found %d unique tags plus %d repeats", infile2, kept, dropped)
                                return err
                        }
                } else if fastaFilenameRe.MatchString(infile) {
                        todo <- func() error {
                                defer phases.Done()
                                defer phases.Done()
-                               log.Printf("%s starting", infile)
+                               log.Printf("%s (reference) starting tiling", infile)
                                defer log.Printf("%s done", infile)
-                               tseqs, stats, err := cmd.tileFasta(tilelib, infile)
+                               tseqs, stats, err := cmd.tileFasta(tilelib, infile, true)
                                allstats[idx*2] = stats
                                if err != nil {
                                        return err
@@ -357,7 +383,17 @@ func (cmd *importer) tileInputs(tilelib *tileLibrary, infiles []string) error {
                                for _, tseq := range tseqs {
                                        totlen += len(tseq)
                                }
-                               log.Printf("%s tiled %d seqs, total len %d", infile, len(tseqs), totlen)
+                               log.Printf("%s (reference) tiled %d seqs, total len %d", infile, len(tseqs), totlen)
+
+                               if cmd.retainAfterEncoding {
+                                       tilelib.mtx.Lock()
+                                       if tilelib.refseqs == nil {
+                                               tilelib.refseqs = map[string]map[string][]tileLibRef{}
+                                       }
+                                       tilelib.refseqs[infile] = tseqs
+                                       tilelib.mtx.Unlock()
+                               }
+
                                return cmd.encoder.Encode(LibraryEntry{
                                        CompactSequences: []CompactSequence{{Name: infile, TileSequences: tseqs}},
                                })
@@ -389,8 +425,9 @@ func (cmd *importer) tileInputs(tilelib *tileLibrary, infiles []string) error {
                        if len(errs) > 0 {
                                return
                        }
+                       variants := flatten(variants)
                        err := cmd.encoder.Encode(LibraryEntry{
-                               CompactGenomes: []CompactGenome{{Name: infile, Variants: flatten(variants)}},
+                               CompactGenomes: []CompactGenome{{Name: infile, Variants: variants}},
                        })
                        if err != nil {
                                select {
@@ -398,12 +435,20 @@ func (cmd *importer) tileInputs(tilelib *tileLibrary, infiles []string) error {
                                default:
                                }
                        }
+                       if cmd.retainAfterEncoding {
+                               tilelib.mtx.Lock()
+                               if tilelib.compactGenomes == nil {
+                                       tilelib.compactGenomes = make(map[string][]tileVariantID)
+                               }
+                               tilelib.compactGenomes[infile] = variants
+                               tilelib.mtx.Unlock()
+                       }
                }()
        }
        go close(todo)
        var tileJobs sync.WaitGroup
        var running int64
-       for i := 0; i < runtime.NumCPU()*9/8+1; i++ {
+       for i := 0; i < runtime.GOMAXPROCS(-1); i++ {
                tileJobs.Add(1)
                atomic.AddInt64(&running, 1)
                go func() {
@@ -430,6 +475,13 @@ func (cmd *importer) tileInputs(tilelib *tileLibrary, infiles []string) error {
                }()
        }
        tileJobs.Wait()
+       if len(errs) > 0 {
+               // Must not wait on encodeJobs in this case. If the
+               // tileJobs goroutines exited early, some funcs in
+               // todo haven't been called, so the corresponding
+               // encodeJobs will wait forever.
+               return <-errs
+       }
        encodeJobs.Wait()
 
        go close(errs)
@@ -488,7 +540,7 @@ func (cmd *importer) tileGVCF(tilelib *tileLibrary, infile string, phase int) (t
                return
        }
        defer consensus.Wait()
-       tileseq, stats, err = tilelib.TileFasta(fmt.Sprintf("%s phase %d", infile, phase+1), stdout, cmd.matchChromosome)
+       tileseq, stats, err = tilelib.TileFasta(fmt.Sprintf("%s phase %d", infile, phase+1), stdout, cmd.matchChromosome, false)
        if err != nil {
                return
        }
@@ -521,73 +573,3 @@ func flatten(variants [][]tileVariantID) []tileVariantID {
        }
        return flat
 }
-
-type importstatsplot struct{}
-
-func (cmd *importstatsplot) RunCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
-       err := cmd.Plot(stdin, stdout)
-       if err != nil {
-               log.Errorf("%s", err)
-               return 1
-       }
-       return 0
-}
-
-func (cmd *importstatsplot) Plot(stdin io.Reader, stdout io.Writer) error {
-       var stats []importStats
-       err := json.NewDecoder(stdin).Decode(&stats)
-       if err != nil {
-               return err
-       }
-
-       p, err := plot.New()
-       if err != nil {
-               return err
-       }
-       p.Title.Text = "coverage preserved by import (excl X<0.65)"
-       p.X.Label.Text = "input base calls ÷ sequence length"
-       p.Y.Label.Text = "output base calls ÷ input base calls"
-       p.Add(plotter.NewGrid())
-
-       data := map[string]plotter.XYs{}
-       for _, stat := range stats {
-               data[stat.InputLabel] = append(data[stat.InputLabel], plotter.XY{
-                       X: float64(stat.InputCoverage) / float64(stat.InputLength),
-                       Y: float64(stat.TileCoverage) / float64(stat.InputCoverage),
-               })
-       }
-
-       labels := []string{}
-       for label := range data {
-               labels = append(labels, label)
-       }
-       sort.Strings(labels)
-       palette, err := colorful.SoftPalette(len(labels))
-       if err != nil {
-               return err
-       }
-       nextInPalette := 0
-       for idx, label := range labels {
-               s, err := plotter.NewScatter(data[label])
-               if err != nil {
-                       return err
-               }
-               s.GlyphStyle.Color = palette[idx]
-               s.GlyphStyle.Radius = vg.Millimeter / 2
-               s.GlyphStyle.Shape = draw.CrossGlyph{}
-               nextInPalette += 7
-               p.Add(s)
-               if false {
-                       p.Legend.Add(label, s)
-               }
-       }
-       p.X.Min = 0.65
-       p.X.Max = 1
-
-       w, err := p.WriterTo(8*vg.Inch, 6*vg.Inch, "svg")
-       if err != nil {
-               return err
-       }
-       _, err = w.WriteTo(stdout)
-       return err
-}