25 "git.arvados.org/arvados.git/sdk/go/arvados"
26 log "github.com/sirupsen/logrus"
29 type importer struct {
42 func (cmd *importer) RunCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
46 fmt.Fprintf(stderr, "%s\n", err)
49 flags := flag.NewFlagSet("", flag.ContinueOnError)
50 flags.SetOutput(stderr)
51 flags.StringVar(&cmd.tagLibraryFile, "tag-library", "", "tag library fasta `file`")
52 flags.StringVar(&cmd.refFile, "ref", "", "reference fasta `file`")
53 flags.StringVar(&cmd.outputFile, "o", "-", "output `file`")
54 flags.StringVar(&cmd.projectUUID, "project", "", "project `UUID` for output data")
55 flags.BoolVar(&cmd.runLocal, "local", false, "run on local host (default: run in an arvados container)")
56 flags.BoolVar(&cmd.skipOOO, "skip-ooo", false, "skip out-of-order tags")
57 flags.BoolVar(&cmd.outputTiles, "output-tiles", false, "include tile variant sequences in output file")
58 flags.BoolVar(&cmd.includeNoCalls, "include-no-calls", false, "treat tiles with no-calls as regular tiles")
59 flags.StringVar(&cmd.outputStats, "output-stats", "", "output stats to `file` (json)")
60 priority := flags.Int("priority", 500, "container request priority")
61 pprof := flags.String("pprof", "", "serve Go profile data at http://`[addr]:port`")
62 loglevel := flags.String("loglevel", "info", "logging threshold (trace, debug, info, warn, error, fatal, or panic)")
63 err = flags.Parse(args)
64 if err == flag.ErrHelp {
67 } else if err != nil {
69 } else if cmd.tagLibraryFile == "" {
70 fmt.Fprintln(os.Stderr, "cannot import without -tag-library argument")
72 } else if flags.NArg() == 0 {
79 log.Println(http.ListenAndServe(*pprof, nil))
83 lvl, err := log.ParseLevel(*loglevel)
90 runner := arvadosContainerRunner{
91 Name: "lightning import",
92 Client: arvados.NewClientFromEnv(),
93 ProjectUUID: cmd.projectUUID,
98 err = runner.TranslatePaths(&cmd.tagLibraryFile, &cmd.refFile, &cmd.outputFile)
102 inputs := flags.Args()
103 for i := range inputs {
104 err = runner.TranslatePaths(&inputs[i])
109 if cmd.outputFile == "-" {
110 cmd.outputFile = "/mnt/output/library.gob"
112 // Not yet implemented, but this should write
113 // the collection to an existing collection,
114 // possibly even an in-place update.
115 err = errors.New("cannot specify output file in container mode: not implemented")
118 runner.Args = append([]string{"import",
120 "-loglevel=" + *loglevel,
121 fmt.Sprintf("-skip-ooo=%v", cmd.skipOOO),
122 fmt.Sprintf("-output-tiles=%v", cmd.outputTiles),
123 fmt.Sprintf("-include-no-calls=%v", cmd.includeNoCalls),
124 "-output-stats", "/mnt/output/stats.json",
125 "-tag-library", cmd.tagLibraryFile,
127 "-o", cmd.outputFile,
130 output, err = runner.Run()
134 fmt.Fprintln(stdout, output+"/library.gob")
138 infiles, err := listInputFiles(flags.Args())
143 taglib, err := cmd.loadTagLibrary()
148 var output io.WriteCloser
149 if cmd.outputFile == "-" {
150 output = nopCloser{stdout}
152 output, err = os.OpenFile(cmd.outputFile, os.O_CREATE|os.O_WRONLY, 0777)
158 bufw := bufio.NewWriter(output)
159 cmd.encoder = gob.NewEncoder(bufw)
161 tilelib := &tileLibrary{taglib: taglib, includeNoCalls: cmd.includeNoCalls, skipOOO: cmd.skipOOO}
163 cmd.encoder.Encode(LibraryEntry{TagSet: taglib.Tags()})
164 tilelib.encoder = cmd.encoder
167 for range time.Tick(10 * time.Minute) {
168 log.Printf("tilelib.Len() == %d", tilelib.Len())
172 err = cmd.tileInputs(tilelib, infiles)
187 func (cmd *importer) tileFasta(tilelib *tileLibrary, infile string) (tileSeq, []importStats, error) {
188 var input io.ReadCloser
189 input, err := os.Open(infile)
194 if strings.HasSuffix(infile, ".gz") {
195 input, err = gzip.NewReader(input)
201 return tilelib.TileFasta(infile, input)
204 func (cmd *importer) loadTagLibrary() (*tagLibrary, error) {
205 log.Printf("tag library %s load starting", cmd.tagLibraryFile)
206 f, err := os.Open(cmd.tagLibraryFile)
211 var rdr io.ReadCloser = f
212 if strings.HasSuffix(cmd.tagLibraryFile, ".gz") {
213 rdr, err = gzip.NewReader(f)
215 return nil, fmt.Errorf("%s: gzip: %s", cmd.tagLibraryFile, err)
219 var taglib tagLibrary
220 err = taglib.Load(rdr)
224 if taglib.Len() < 1 {
225 return nil, fmt.Errorf("cannot tile: tag library is empty")
227 log.Printf("tag library %s load done", cmd.tagLibraryFile)
232 vcfFilenameRe = regexp.MustCompile(`\.vcf(\.gz)?$`)
233 fasta1FilenameRe = regexp.MustCompile(`\.1\.fa(sta)?(\.gz)?$`)
234 fasta2FilenameRe = regexp.MustCompile(`\.2\.fa(sta)?(\.gz)?$`)
235 fastaFilenameRe = regexp.MustCompile(`\.fa(sta)?(\.gz)?$`)
238 func listInputFiles(paths []string) (files []string, err error) {
239 for _, path := range paths {
240 if fi, err := os.Stat(path); err != nil {
241 return nil, fmt.Errorf("%s: stat failed: %s", path, err)
242 } else if !fi.IsDir() {
243 if !fasta2FilenameRe.MatchString(path) {
244 files = append(files, path)
248 d, err := os.Open(path)
250 return nil, fmt.Errorf("%s: open failed: %s", path, err)
253 names, err := d.Readdirnames(0)
255 return nil, fmt.Errorf("%s: readdir failed: %s", path, err)
258 for _, name := range names {
259 if vcfFilenameRe.MatchString(name) {
260 files = append(files, filepath.Join(path, name))
261 } else if fastaFilenameRe.MatchString(name) && !fasta2FilenameRe.MatchString(name) {
262 files = append(files, filepath.Join(path, name))
267 for _, file := range files {
268 if fastaFilenameRe.MatchString(file) {
270 } else if vcfFilenameRe.MatchString(file) {
271 if _, err := os.Stat(file + ".csi"); err == nil {
273 } else if _, err = os.Stat(file + ".tbi"); err == nil {
276 return nil, fmt.Errorf("%s: cannot read without .tbi or .csi index file", file)
279 return nil, fmt.Errorf("don't know how to handle filename %s", file)
285 func (cmd *importer) tileInputs(tilelib *tileLibrary, infiles []string) error {
286 starttime := time.Now()
287 errs := make(chan error, 1)
288 todo := make(chan func() error, len(infiles)*2)
289 allstats := make([][]importStats, len(infiles)*2)
290 var encodeJobs sync.WaitGroup
291 for idx, infile := range infiles {
292 idx, infile := idx, infile
293 var phases sync.WaitGroup
295 variants := make([][]tileVariantID, 2)
296 if fasta1FilenameRe.MatchString(infile) {
297 todo <- func() error {
299 log.Printf("%s starting", infile)
300 defer log.Printf("%s done", infile)
301 tseqs, stats, err := cmd.tileFasta(tilelib, infile)
302 allstats[idx*2] = stats
303 var kept, dropped int
304 variants[0], kept, dropped = tseqs.Variants()
305 log.Printf("%s found %d unique tags plus %d repeats", infile, kept, dropped)
308 infile2 := fasta1FilenameRe.ReplaceAllString(infile, `.2.fa$1$2`)
309 todo <- func() error {
311 log.Printf("%s starting", infile2)
312 defer log.Printf("%s done", infile2)
313 tseqs, stats, err := cmd.tileFasta(tilelib, infile2)
314 allstats[idx*2+1] = stats
315 var kept, dropped int
316 variants[1], kept, dropped = tseqs.Variants()
317 log.Printf("%s found %d unique tags plus %d repeats", infile2, kept, dropped)
321 } else if fastaFilenameRe.MatchString(infile) {
322 todo <- func() error {
325 log.Printf("%s starting", infile)
326 defer log.Printf("%s done", infile)
327 tseqs, stats, err := cmd.tileFasta(tilelib, infile)
328 allstats[idx*2] = stats
333 for _, tseq := range tseqs {
336 log.Printf("%s tiled %d seqs, total len %d", infile, len(tseqs), totlen)
337 return cmd.encoder.Encode(LibraryEntry{
338 CompactSequences: []CompactSequence{{Name: infile, TileSequences: tseqs}},
341 // Don't write out a CompactGenomes entry
343 } else if vcfFilenameRe.MatchString(infile) {
344 for phase := 0; phase < 2; phase++ {
346 todo <- func() error {
348 log.Printf("%s phase %d starting", infile, phase+1)
349 defer log.Printf("%s phase %d done", infile, phase+1)
350 tseqs, stats, err := cmd.tileGVCF(tilelib, infile, phase)
351 allstats[idx*2] = stats
352 var kept, dropped int
353 variants[phase], kept, dropped = tseqs.Variants()
354 log.Printf("%s phase %d found %d unique tags plus %d repeats", infile, phase+1, kept, dropped)
359 panic(fmt.Sprintf("bug: unhandled filename %q", infile))
363 defer encodeJobs.Done()
368 err := cmd.encoder.Encode(LibraryEntry{
369 CompactGenomes: []CompactGenome{{Name: infile, Variants: flatten(variants)}},
380 var tileJobs sync.WaitGroup
382 for i := 0; i < runtime.NumCPU()*9/8+1; i++ {
384 atomic.AddInt64(&running, 1)
386 defer tileJobs.Done()
387 defer atomic.AddInt64(&running, -1)
388 for fn := range todo {
399 remain := len(todo) + int(atomic.LoadInt64(&running)) - 1
400 if remain < cap(todo) {
401 ttl := time.Now().Sub(starttime) * time.Duration(remain) / time.Duration(cap(todo)-remain)
402 eta := time.Now().Add(ttl)
403 log.Printf("progress %d/%d, eta %v (%v)", cap(todo)-remain, cap(todo), eta, ttl)
417 if cmd.outputStats != "" {
418 f, err := os.OpenFile(cmd.outputStats, os.O_CREATE|os.O_WRONLY, 0666)
422 var flatstats []importStats
423 for _, stats := range allstats {
424 flatstats = append(flatstats, stats...)
426 err = json.NewEncoder(f).Encode(flatstats)
435 func (cmd *importer) tileGVCF(tilelib *tileLibrary, infile string, phase int) (tileseq tileSeq, stats []importStats, err error) {
436 if cmd.refFile == "" {
437 err = errors.New("cannot import vcf: reference data (-ref) not specified")
440 args := []string{"bcftools", "consensus", "--fasta-ref", cmd.refFile, "-H", fmt.Sprint(phase + 1), infile}
441 indexsuffix := ".tbi"
442 if _, err := os.Stat(infile + ".csi"); err == nil {
445 if out, err := exec.Command("docker", "image", "ls", "-q", "lightning-runtime").Output(); err == nil && len(out) > 0 {
446 args = append([]string{
447 "docker", "run", "--rm",
449 "--volume=" + infile + ":" + infile + ":ro",
450 "--volume=" + infile + indexsuffix + ":" + infile + indexsuffix + ":ro",
451 "--volume=" + cmd.refFile + ":" + cmd.refFile + ":ro",
455 consensus := exec.Command(args[0], args[1:]...)
456 consensus.Stderr = os.Stderr
457 stdout, err := consensus.StdoutPipe()
462 err = consensus.Start()
466 defer consensus.Wait()
467 tileseq, stats, err = tilelib.TileFasta(fmt.Sprintf("%s phase %d", infile, phase+1), stdout)
475 err = consensus.Wait()
477 err = fmt.Errorf("%s phase %d: bcftools: %s", infile, phase, err)
483 func flatten(variants [][]tileVariantID) []tileVariantID {
485 for _, v := range variants {
490 flat := make([]tileVariantID, ntags*2)
491 for i := 0; i < ntags; i++ {
492 for hap := 0; hap < 2; hap++ {
493 if i < len(variants[hap]) {
494 flat[i*2+hap] = variants[hap][i]