X-Git-Url: https://git.arvados.org/lightning.git/blobdiff_plain/40c17a78f9cdc073f494052a3421d43ba4a9d441..HEAD:/import.go

diff --git a/import.go b/import.go
index 14b3667ed7..20058dc8a9 100644
--- 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"
-	"log"
+	"io/ioutil"
 	"net/http"
 	_ "net/http/pprof"
 	"os"
@@ -22,16 +28,26 @@ import (
 	"sync/atomic"
 	"time"
 
-	"git.arvados.org/arvados.git/sdk/go/arvados"
+	"github.com/klauspost/pgzip"
+	log "github.com/sirupsen/logrus"
 )
 
 type importer struct {
-	tagLibraryFile string
-	refFile        string
-	outputFile     string
-	projectUUID    string
-	runLocal       bool
-	encoder        *gob.Encoder
+	tagLibraryFile      string
+	refFile             string
+	outputFile          string
+	projectUUID         string
+	loglevel            string
+	priority            int
+	runLocal            bool
+	skipOOO             bool
+	outputTiles         bool
+	saveIncompleteTiles bool
+	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 {
@@ -45,10 +61,18 @@ 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)")
+	flags.BoolVar(&cmd.skipOOO, "skip-ooo", false, "skip out-of-order tags")
+	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`")
+	flags.IntVar(&cmd.priority, "priority", 500, "container request priority")
 	pprof := flags.String("pprof", "", "serve Go profile data at http://`[addr]:port`")
+	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
@@ -69,36 +93,19 @@ func (cmd *importer) RunCommand(prog string, args []string, stdin io.Reader, std
 		}()
 	}
 
+	lvl, err := log.ParseLevel(cmd.loglevel)
+	if err != nil {
+		return 2
+	}
+	log.SetLevel(lvl)
+
+	cmd.matchChromosome, err = regexp.Compile(*matchChromosome)
+	if err != nil {
+		return 1
+	}
+
 	if !cmd.runLocal {
-		runner := arvadosContainerRunner{
-			Name:        "lightning import",
-			Client:      arvados.NewClientFromEnv(),
-			ProjectUUID: cmd.projectUUID,
-			RAM:         30000000000,
-			VCPUs:       16,
-		}
-		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...)
-		err = runner.Run()
+		err = cmd.runBatches(stdout, flags.Args())
 		if err != nil {
 			return 1
 		}
@@ -109,30 +116,42 @@ func (cmd *importer) RunCommand(prog string, args []string, stdin io.Reader, std
 	if err != nil {
 		return 1
 	}
+	infiles = cmd.batchArgs.Slice(infiles)
 
-	tilelib, err := cmd.loadTileLibrary()
+	taglib, err := cmd.loadTagLibrary()
 	if err != nil {
 		return 1
 	}
-	go func() {
-		for range time.Tick(10 * time.Second) {
-			log.Printf("tilelib.Len() == %d", tilelib.Len())
-		}
-	}()
 
-	var output io.WriteCloser
-	if cmd.outputFile == "" {
-		output = nopCloser{stdout}
+	var outw, outf io.WriteCloser
+	if cmd.outputFile == "-" {
+		outw = nopCloser{stdout}
 	} else {
-		output, err = os.OpenFile(cmd.outputFile, os.O_CREATE|os.O_WRONLY, 0777)
+		outf, err = os.OpenFile(cmd.outputFile, os.O_CREATE|os.O_WRONLY, 0777)
 		if err != nil {
 			return 1
 		}
-		defer output.Close()
+		defer outf.Close()
+		if strings.HasSuffix(cmd.outputFile, ".gz") {
+			outw = pgzip.NewWriter(outf)
+		} else {
+			outw = outf
+		}
 	}
-	bufw := bufio.NewWriter(output)
+	bufw := bufio.NewWriterSize(outw, 64*1024*1024)
 	cmd.encoder = gob.NewEncoder(bufw)
 
+	tilelib := &tileLibrary{taglib: taglib, retainNoCalls: cmd.saveIncompleteTiles, skipOOO: cmd.skipOOO}
+	if cmd.outputTiles {
+		cmd.encoder.Encode(LibraryEntry{TagSet: taglib.Tags()})
+		tilelib.encoder = cmd.encoder
+	}
+	go func() {
+		for range time.Tick(10 * time.Minute) {
+			log.Printf("tilelib.Len() == %d", tilelib.Len())
+		}
+	}()
+
 	err = cmd.tileInputs(tilelib, infiles)
 	if err != nil {
 		return 1
@@ -141,40 +160,108 @@ func (cmd *importer) RunCommand(prog string, args []string, stdin io.Reader, std
 	if err != nil {
 		return 1
 	}
-	err = output.Close()
+	err = outw.Close()
 	if err != nil {
 		return 1
 	}
+	if outf != nil && outf != outw {
+		err = outf.Close()
+		if err != nil {
+			return 1
+		}
+	}
 	return 0
 }
 
-func (cmd *importer) tileFasta(tilelib *tileLibrary, infile string) (tileSeq, 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, err
+		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, err
+			return nil, nil, err
 		}
 		defer input.Close()
 	}
-	return tilelib.TileFasta(infile, input)
+	return tilelib.TileFasta(infile, input, cmd.matchChromosome, isRef)
 }
 
-func (cmd *importer) loadTileLibrary() (*tileLibrary, error) {
+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)
 		}
@@ -189,15 +276,22 @@ func (cmd *importer) loadTileLibrary() (*tileLibrary, error) {
 		return nil, fmt.Errorf("cannot tile: tag library is empty")
 	}
 	log.Printf("tag library %s load done", cmd.tagLibraryFile)
-	return &tileLibrary{taglib: &taglib}, nil
+	return &taglib, nil
 }
 
+var (
+	vcfFilenameRe    = regexp.MustCompile(`\.vcf(\.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)?$`)
+)
+
 func listInputFiles(paths []string) (files []string, err error) {
 	for _, path := range paths {
 		if fi, err := os.Stat(path); err != nil {
 			return nil, fmt.Errorf("%s: stat failed: %s", path, err)
 		} else if !fi.IsDir() {
-			if !strings.HasSuffix(path, ".2.fasta") || strings.HasSuffix(path, ".2.fasta.gz") {
+			if !fasta2FilenameRe.MatchString(path) {
 				files = append(files, path)
 			}
 			continue
@@ -213,23 +307,27 @@ func listInputFiles(paths []string) (files []string, err error) {
 		}
 		sort.Strings(names)
 		for _, name := range names {
-			if strings.HasSuffix(name, ".vcf") || strings.HasSuffix(name, ".vcf.gz") {
+			if vcfFilenameRe.MatchString(name) {
 				files = append(files, filepath.Join(path, name))
-			} else if strings.HasSuffix(name, ".1.fasta") || strings.HasSuffix(name, ".1.fasta.gz") {
+			} else if fastaFilenameRe.MatchString(name) && !fasta2FilenameRe.MatchString(name) {
 				files = append(files, filepath.Join(path, name))
 			}
 		}
 		d.Close()
 	}
 	for _, file := range files {
-		if strings.HasSuffix(file, ".1.fasta") || strings.HasSuffix(file, ".1.fasta.gz") {
-			continue
-		} else if _, err := os.Stat(file + ".csi"); err == nil {
-			continue
-		} else if _, err = os.Stat(file + ".tbi"); err == nil {
+		if fastaFilenameRe.MatchString(file) {
 			continue
+		} else if vcfFilenameRe.MatchString(file) {
+			if _, err := os.Stat(file + ".csi"); err == nil {
+				continue
+			} else if _, err = os.Stat(file + ".tbi"); err == nil {
+				continue
+			} else {
+				return nil, fmt.Errorf("%s: cannot read without .tbi or .csi index file", file)
+			}
 		} else {
-			return nil, fmt.Errorf("%s: cannot read without .tbi or .csi index file", file)
+			return nil, fmt.Errorf("don't know how to handle filename %s", file)
 		}
 	}
 	return
@@ -239,42 +337,86 @@ func (cmd *importer) tileInputs(tilelib *tileLibrary, infiles []string) error {
 	starttime := time.Now()
 	errs := make(chan error, 1)
 	todo := make(chan func() error, len(infiles)*2)
+	allstats := make([][]importStats, len(infiles)*2)
 	var encodeJobs sync.WaitGroup
-	for _, infile := range infiles {
-		infile := infile
+	for idx, infile := range infiles {
+		idx, infile := idx, infile
 		var phases sync.WaitGroup
 		phases.Add(2)
 		variants := make([][]tileVariantID, 2)
-		if strings.HasSuffix(infile, ".1.fasta") || strings.HasSuffix(infile, ".1.fasta.gz") {
+		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, err := cmd.tileFasta(tilelib, infile)
-				variants[0] = tseqs.Variants()
+				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 (sample.1) found %d unique tags plus %d repeats", infile, kept, dropped)
 				return err
 			}
-			infile2 := regexp.MustCompile(`\.1\.fasta(\.gz)?$`).ReplaceAllString(infile, `.2.fasta$1`)
+			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, err := cmd.tileFasta(tilelib, infile2)
-				variants[1] = tseqs.Variants()
+				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 (sample.2) found %d unique tags plus %d repeats", infile2, kept, dropped)
 				return err
 			}
-		} else {
+		} else if fastaFilenameRe.MatchString(infile) {
+			todo <- func() error {
+				defer phases.Done()
+				defer phases.Done()
+				log.Printf("%s (reference) starting tiling", infile)
+				defer log.Printf("%s done", infile)
+				tseqs, stats, err := cmd.tileFasta(tilelib, infile, true)
+				allstats[idx*2] = stats
+				if err != nil {
+					return err
+				}
+				totlen := 0
+				for _, tseq := range tseqs {
+					totlen += len(tseq)
+				}
+				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}},
+				})
+			}
+			// Don't write out a CompactGenomes entry
+			continue
+		} else if vcfFilenameRe.MatchString(infile) {
 			for phase := 0; phase < 2; phase++ {
 				phase := phase
 				todo <- func() error {
 					defer phases.Done()
 					log.Printf("%s phase %d starting", infile, phase+1)
 					defer log.Printf("%s phase %d done", infile, phase+1)
-					tseqs, err := cmd.tileGVCF(tilelib, infile, phase)
-					variants[phase] = tseqs.Variants()
+					tseqs, stats, err := cmd.tileGVCF(tilelib, infile, phase)
+					allstats[idx*2] = stats
+					var kept, dropped int
+					variants[phase], kept, dropped = tseqs.Variants()
+					log.Printf("%s phase %d found %d unique tags plus %d repeats", infile, phase+1, kept, dropped)
 					return err
 				}
 			}
+		} else {
+			panic(fmt.Sprintf("bug: unhandled filename %q", infile))
 		}
 		encodeJobs.Add(1)
 		go func() {
@@ -283,17 +425,9 @@ func (cmd *importer) tileInputs(tilelib *tileLibrary, infiles []string) error {
 			if len(errs) > 0 {
 				return
 			}
-			ntags := len(variants[0])
-			if ntags < len(variants[1]) {
-				ntags = len(variants[1])
-			}
-			flat := make([]tileVariantID, ntags*2)
-			for i := 0; i < ntags; i++ {
-				flat[i*2] = variants[0][i]
-				flat[i*2+1] = variants[1][i]
-			}
+			variants := flatten(variants)
 			err := cmd.encoder.Encode(LibraryEntry{
-				CompactGenomes: []CompactGenome{{Name: infile, Variants: flat}},
+				CompactGenomes: []CompactGenome{{Name: infile, Variants: variants}},
 			})
 			if err != nil {
 				select {
@@ -301,13 +435,22 @@ 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
-	running := int64(runtime.NumCPU())
-	for i := 0; i < runtime.NumCPU(); i++ {
+	var running int64
+	for i := 0; i < runtime.GOMAXPROCS(-1); i++ {
 		tileJobs.Add(1)
+		atomic.AddInt64(&running, 1)
 		go func() {
 			defer tileJobs.Done()
 			defer atomic.AddInt64(&running, -1)
@@ -323,19 +466,49 @@ func (cmd *importer) tileInputs(tilelib *tileLibrary, infiles []string) error {
 					}
 				}
 				remain := len(todo) + int(atomic.LoadInt64(&running)) - 1
-				ttl := time.Now().Sub(starttime) * time.Duration(remain) / time.Duration(cap(todo)-remain)
-				eta := time.Now().Add(ttl)
-				log.Printf("progress %d/%d, eta %v (%v)", cap(todo)-remain, cap(todo), eta, ttl)
+				if remain < cap(todo) {
+					ttl := time.Now().Sub(starttime) * time.Duration(remain) / time.Duration(cap(todo)-remain)
+					eta := time.Now().Add(ttl)
+					log.Printf("progress %d/%d, eta %v (%v)", cap(todo)-remain, cap(todo), eta, ttl)
+				}
 			}
 		}()
 	}
 	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)
-	return <-errs
+	err := <-errs
+	if err != nil {
+		return err
+	}
+
+	if cmd.outputStats != "" {
+		f, err := os.OpenFile(cmd.outputStats, os.O_CREATE|os.O_WRONLY, 0666)
+		if err != nil {
+			return err
+		}
+		var flatstats []importStats
+		for _, stats := range allstats {
+			flatstats = append(flatstats, stats...)
+		}
+		err = json.NewEncoder(f).Encode(flatstats)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
 }
 
-func (cmd *importer) tileGVCF(tilelib *tileLibrary, infile string, phase int) (tileseq tileSeq, err error) {
+func (cmd *importer) tileGVCF(tilelib *tileLibrary, infile string, phase int) (tileseq tileSeq, stats []importStats, err error) {
 	if cmd.refFile == "" {
 		err = errors.New("cannot import vcf: reference data (-ref) not specified")
 		return
@@ -367,7 +540,7 @@ func (cmd *importer) tileGVCF(tilelib *tileLibrary, infile string, phase int) (t
 		return
 	}
 	defer consensus.Wait()
-	tileseq, err = tilelib.TileFasta(fmt.Sprintf("%s phase %d", infile, phase+1), stdout)
+	tileseq, stats, err = tilelib.TileFasta(fmt.Sprintf("%s phase %d", infile, phase+1), stdout, cmd.matchChromosome, false)
 	if err != nil {
 		return
 	}
@@ -382,3 +555,21 @@ func (cmd *importer) tileGVCF(tilelib *tileLibrary, infile string, phase int) (t
 	}
 	return
 }
+
+func flatten(variants [][]tileVariantID) []tileVariantID {
+	ntags := 0
+	for _, v := range variants {
+		if ntags < len(v) {
+			ntags = len(v)
+		}
+	}
+	flat := make([]tileVariantID, ntags*2)
+	for i := 0; i < ntags; i++ {
+		for hap := 0; hap < 2; hap++ {
+			if i < len(variants[hap]) {
+				flat[i*2+hap] = variants[hap][i]
+			}
+		}
+	}
+	return flat
+}