import.go

   1 package main
   2
   3 import (
   4         "bufio"
   5         "compress/gzip"
   6         "encoding/gob"
   7         "encoding/json"
   8         "errors"
   9         "flag"
  10         "fmt"
  11         "io"
  12         "net/http"
  13         _ "net/http/pprof"
  14         "os"
  15         "os/exec"
  16         "path/filepath"
  17         "regexp"
  18         "runtime"
  19         "sort"
  20         "strings"
  21         "sync"
  22         "sync/atomic"
  23         "time"
  24
  25         "git.arvados.org/arvados.git/sdk/go/arvados"
  26         "github.com/lucasb-eyer/go-colorful"
  27         log "github.com/sirupsen/logrus"
  28         "gonum.org/v1/plot"
  29         "gonum.org/v1/plot/plotter"
  30         "gonum.org/v1/plot/vg"
  31         "gonum.org/v1/plot/vg/draw"
  32 )
  33
  34 type importer struct {
  35         tagLibraryFile      string
  36         refFile             string
  37         outputFile          string
  38         projectUUID         string
  39         runLocal            bool
  40         skipOOO             bool
  41         outputTiles         bool
  42         saveIncompleteTiles bool
  43         outputStats         string
  44         encoder             *gob.Encoder
  45 }
  46
  47 func (cmd *importer) RunCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
  48         var err error
  49         defer func() {
  50                 if err != nil {
  51                         fmt.Fprintf(stderr, "%s\n", err)
  52                 }
  53         }()
  54         flags := flag.NewFlagSet("", flag.ContinueOnError)
  55         flags.SetOutput(stderr)
  56         flags.StringVar(&cmd.tagLibraryFile, "tag-library", "", "tag library fasta `file`")
  57         flags.StringVar(&cmd.refFile, "ref", "", "reference fasta `file`")
  58         flags.StringVar(&cmd.outputFile, "o", "-", "output `file`")
  59         flags.StringVar(&cmd.projectUUID, "project", "", "project `UUID` for output data")
  60         flags.BoolVar(&cmd.runLocal, "local", false, "run on local host (default: run in an arvados container)")
  61         flags.BoolVar(&cmd.skipOOO, "skip-ooo", false, "skip out-of-order tags")
  62         flags.BoolVar(&cmd.outputTiles, "output-tiles", false, "include tile variant sequences in output file")
  63         flags.BoolVar(&cmd.saveIncompleteTiles, "save-incomplete-tiles", false, "treat tiles with no-calls as regular tiles")
  64         flags.StringVar(&cmd.outputStats, "output-stats", "", "output stats to `file` (json)")
  65         priority := flags.Int("priority", 500, "container request priority")
  66         pprof := flags.String("pprof", "", "serve Go profile data at http://`[addr]:port`")
  67         loglevel := flags.String("loglevel", "info", "logging threshold (trace, debug, info, warn, error, fatal, or panic)")
  68         err = flags.Parse(args)
  69         if err == flag.ErrHelp {
  70                 err = nil
  71                 return 0
  72         } else if err != nil {
  73                 return 2
  74         } else if cmd.tagLibraryFile == "" {
  75                 fmt.Fprintln(os.Stderr, "cannot import without -tag-library argument")
  76                 return 2
  77         } else if flags.NArg() == 0 {
  78                 flags.Usage()
  79                 return 2
  80         }
  81
  82         if *pprof != "" {
  83                 go func() {
  84                         log.Println(http.ListenAndServe(*pprof, nil))
  85                 }()
  86         }
  87
  88         lvl, err := log.ParseLevel(*loglevel)
  89         if err != nil {
  90                 return 2
  91         }
  92         log.SetLevel(lvl)
  93
  94         if !cmd.runLocal {
  95                 runner := arvadosContainerRunner{
  96                         Name:        "lightning import",
  97                         Client:      arvados.NewClientFromEnv(),
  98                         ProjectUUID: cmd.projectUUID,
  99                         RAM:         80000000000,
 100                         VCPUs:       32,
 101                         Priority:    *priority,
 102                 }
 103                 err = runner.TranslatePaths(&cmd.tagLibraryFile, &cmd.refFile, &cmd.outputFile)
 104                 if err != nil {
 105                         return 1
 106                 }
 107                 inputs := flags.Args()
 108                 for i := range inputs {
 109                         err = runner.TranslatePaths(&inputs[i])
 110                         if err != nil {
 111                                 return 1
 112                         }
 113                 }
 114                 if cmd.outputFile == "-" {
 115                         cmd.outputFile = "/mnt/output/library.gob.gz"
 116                 } else {
 117                         // Not yet implemented, but this should write
 118                         // the collection to an existing collection,
 119                         // possibly even an in-place update.
 120                         err = errors.New("cannot specify output file in container mode: not implemented")
 121                         return 1
 122                 }
 123                 runner.Args = append([]string{"import",
 124                         "-local=true",
 125                         "-loglevel=" + *loglevel,
 126                         fmt.Sprintf("-skip-ooo=%v", cmd.skipOOO),
 127                         fmt.Sprintf("-output-tiles=%v", cmd.outputTiles),
 128                         fmt.Sprintf("-save-incomplete-tiles=%v", cmd.saveIncompleteTiles),
 129                         "-output-stats", "/mnt/output/stats.json",
 130                         "-tag-library", cmd.tagLibraryFile,
 131                         "-ref", cmd.refFile,
 132                         "-o", cmd.outputFile,
 133                 }, inputs...)
 134                 var output string
 135                 output, err = runner.Run()
 136                 if err != nil {
 137                         return 1
 138                 }
 139                 fmt.Fprintln(stdout, output+"/library.gob.gz")
 140                 return 0
 141         }
 142
 143         infiles, err := listInputFiles(flags.Args())
 144         if err != nil {
 145                 return 1
 146         }
 147
 148         taglib, err := cmd.loadTagLibrary()
 149         if err != nil {
 150                 return 1
 151         }
 152
 153         var outw, outf io.WriteCloser
 154         if cmd.outputFile == "-" {
 155                 outw = nopCloser{stdout}
 156         } else {
 157                 outf, err = os.OpenFile(cmd.outputFile, os.O_CREATE|os.O_WRONLY, 0777)
 158                 if err != nil {
 159                         return 1
 160                 }
 161                 defer outf.Close()
 162                 if strings.HasSuffix(cmd.outputFile, ".gz") {
 163                         outw = gzip.NewWriter(outf)
 164                 } else {
 165                         outw = outf
 166                 }
 167         }
 168         bufw := bufio.NewWriter(outw)
 169         cmd.encoder = gob.NewEncoder(bufw)
 170
 171         tilelib := &tileLibrary{taglib: taglib, retainNoCalls: cmd.saveIncompleteTiles, skipOOO: cmd.skipOOO}
 172         if cmd.outputTiles {
 173                 cmd.encoder.Encode(LibraryEntry{TagSet: taglib.Tags()})
 174                 tilelib.encoder = cmd.encoder
 175         }
 176         go func() {
 177                 for range time.Tick(10 * time.Minute) {
 178                         log.Printf("tilelib.Len() == %d", tilelib.Len())
 179                 }
 180         }()
 181
 182         err = cmd.tileInputs(tilelib, infiles)
 183         if err != nil {
 184                 return 1
 185         }
 186         err = bufw.Flush()
 187         if err != nil {
 188                 return 1
 189         }
 190         err = outw.Close()
 191         if err != nil {
 192                 return 1
 193         }
 194         if outf != nil && outf != outw {
 195                 err = outf.Close()
 196                 if err != nil {
 197                         return 1
 198                 }
 199         }
 200         return 0
 201 }
 202
 203 func (cmd *importer) tileFasta(tilelib *tileLibrary, infile string) (tileSeq, []importStats, error) {
 204         var input io.ReadCloser
 205         input, err := os.Open(infile)
 206         if err != nil {
 207                 return nil, nil, err
 208         }
 209         defer input.Close()
 210         if strings.HasSuffix(infile, ".gz") {
 211                 input, err = gzip.NewReader(input)
 212                 if err != nil {
 213                         return nil, nil, err
 214                 }
 215                 defer input.Close()
 216         }
 217         return tilelib.TileFasta(infile, input)
 218 }
 219
 220 func (cmd *importer) loadTagLibrary() (*tagLibrary, error) {
 221         log.Printf("tag library %s load starting", cmd.tagLibraryFile)
 222         f, err := os.Open(cmd.tagLibraryFile)
 223         if err != nil {
 224                 return nil, err
 225         }
 226         defer f.Close()
 227         var rdr io.ReadCloser = f
 228         if strings.HasSuffix(cmd.tagLibraryFile, ".gz") {
 229                 rdr, err = gzip.NewReader(f)
 230                 if err != nil {
 231                         return nil, fmt.Errorf("%s: gzip: %s", cmd.tagLibraryFile, err)
 232                 }
 233                 defer rdr.Close()
 234         }
 235         var taglib tagLibrary
 236         err = taglib.Load(rdr)
 237         if err != nil {
 238                 return nil, err
 239         }
 240         if taglib.Len() < 1 {
 241                 return nil, fmt.Errorf("cannot tile: tag library is empty")
 242         }
 243         log.Printf("tag library %s load done", cmd.tagLibraryFile)
 244         return &taglib, nil
 245 }
 246
 247 var (
 248         vcfFilenameRe    = regexp.MustCompile(`\.vcf(\.gz)?$`)
 249         fasta1FilenameRe = regexp.MustCompile(`\.1\.fa(sta)?(\.gz)?$`)
 250         fasta2FilenameRe = regexp.MustCompile(`\.2\.fa(sta)?(\.gz)?$`)
 251         fastaFilenameRe  = regexp.MustCompile(`\.fa(sta)?(\.gz)?$`)
 252 )
 253
 254 func listInputFiles(paths []string) (files []string, err error) {
 255         for _, path := range paths {
 256                 if fi, err := os.Stat(path); err != nil {
 257                         return nil, fmt.Errorf("%s: stat failed: %s", path, err)
 258                 } else if !fi.IsDir() {
 259                         if !fasta2FilenameRe.MatchString(path) {
 260                                 files = append(files, path)
 261                         }
 262                         continue
 263                 }
 264                 d, err := os.Open(path)
 265                 if err != nil {
 266                         return nil, fmt.Errorf("%s: open failed: %s", path, err)
 267                 }
 268                 defer d.Close()
 269                 names, err := d.Readdirnames(0)
 270                 if err != nil {
 271                         return nil, fmt.Errorf("%s: readdir failed: %s", path, err)
 272                 }
 273                 sort.Strings(names)
 274                 for _, name := range names {
 275                         if vcfFilenameRe.MatchString(name) {
 276                                 files = append(files, filepath.Join(path, name))
 277                         } else if fastaFilenameRe.MatchString(name) && !fasta2FilenameRe.MatchString(name) {
 278                                 files = append(files, filepath.Join(path, name))
 279                         }
 280                 }
 281                 d.Close()
 282         }
 283         for _, file := range files {
 284                 if fastaFilenameRe.MatchString(file) {
 285                         continue
 286                 } else if vcfFilenameRe.MatchString(file) {
 287                         if _, err := os.Stat(file + ".csi"); err == nil {
 288                                 continue
 289                         } else if _, err = os.Stat(file + ".tbi"); err == nil {
 290                                 continue
 291                         } else {
 292                                 return nil, fmt.Errorf("%s: cannot read without .tbi or .csi index file", file)
 293                         }
 294                 } else {
 295                         return nil, fmt.Errorf("don't know how to handle filename %s", file)
 296                 }
 297         }
 298         return
 299 }
 300
 301 func (cmd *importer) tileInputs(tilelib *tileLibrary, infiles []string) error {
 302         starttime := time.Now()
 303         errs := make(chan error, 1)
 304         todo := make(chan func() error, len(infiles)*2)
 305         allstats := make([][]importStats, len(infiles)*2)
 306         var encodeJobs sync.WaitGroup
 307         for idx, infile := range infiles {
 308                 idx, infile := idx, infile
 309                 var phases sync.WaitGroup
 310                 phases.Add(2)
 311                 variants := make([][]tileVariantID, 2)
 312                 if fasta1FilenameRe.MatchString(infile) {
 313                         todo <- func() error {
 314                                 defer phases.Done()
 315                                 log.Printf("%s starting", infile)
 316                                 defer log.Printf("%s done", infile)
 317                                 tseqs, stats, err := cmd.tileFasta(tilelib, infile)
 318                                 allstats[idx*2] = stats
 319                                 var kept, dropped int
 320                                 variants[0], kept, dropped = tseqs.Variants()
 321                                 log.Printf("%s found %d unique tags plus %d repeats", infile, kept, dropped)
 322                                 return err
 323                         }
 324                         infile2 := fasta1FilenameRe.ReplaceAllString(infile, `.2.fa$1$2`)
 325                         todo <- func() error {
 326                                 defer phases.Done()
 327                                 log.Printf("%s starting", infile2)
 328                                 defer log.Printf("%s done", infile2)
 329                                 tseqs, stats, err := cmd.tileFasta(tilelib, infile2)
 330                                 allstats[idx*2+1] = stats
 331                                 var kept, dropped int
 332                                 variants[1], kept, dropped = tseqs.Variants()
 333                                 log.Printf("%s found %d unique tags plus %d repeats", infile2, kept, dropped)
 334
 335                                 return err
 336                         }
 337                 } else if fastaFilenameRe.MatchString(infile) {
 338                         todo <- func() error {
 339                                 defer phases.Done()
 340                                 defer phases.Done()
 341                                 log.Printf("%s starting", infile)
 342                                 defer log.Printf("%s done", infile)
 343                                 tseqs, stats, err := cmd.tileFasta(tilelib, infile)
 344                                 allstats[idx*2] = stats
 345                                 if err != nil {
 346                                         return err
 347                                 }
 348                                 totlen := 0
 349                                 for _, tseq := range tseqs {
 350                                         totlen += len(tseq)
 351                                 }
 352                                 log.Printf("%s tiled %d seqs, total len %d", infile, len(tseqs), totlen)
 353                                 return cmd.encoder.Encode(LibraryEntry{
 354                                         CompactSequences: []CompactSequence{{Name: infile, TileSequences: tseqs}},
 355                                 })
 356                         }
 357                         // Don't write out a CompactGenomes entry
 358                         continue
 359                 } else if vcfFilenameRe.MatchString(infile) {
 360                         for phase := 0; phase < 2; phase++ {
 361                                 phase := phase
 362                                 todo <- func() error {
 363                                         defer phases.Done()
 364                                         log.Printf("%s phase %d starting", infile, phase+1)
 365                                         defer log.Printf("%s phase %d done", infile, phase+1)
 366                                         tseqs, stats, err := cmd.tileGVCF(tilelib, infile, phase)
 367                                         allstats[idx*2] = stats
 368                                         var kept, dropped int
 369                                         variants[phase], kept, dropped = tseqs.Variants()
 370                                         log.Printf("%s phase %d found %d unique tags plus %d repeats", infile, phase+1, kept, dropped)
 371                                         return err
 372                                 }
 373                         }
 374                 } else {
 375                         panic(fmt.Sprintf("bug: unhandled filename %q", infile))
 376                 }
 377                 encodeJobs.Add(1)
 378                 go func() {
 379                         defer encodeJobs.Done()
 380                         phases.Wait()
 381                         if len(errs) > 0 {
 382                                 return
 383                         }
 384                         err := cmd.encoder.Encode(LibraryEntry{
 385                                 CompactGenomes: []CompactGenome{{Name: infile, Variants: flatten(variants)}},
 386                         })
 387                         if err != nil {
 388                                 select {
 389                                 case errs <- err:
 390                                 default:
 391                                 }
 392                         }
 393                 }()
 394         }
 395         go close(todo)
 396         var tileJobs sync.WaitGroup
 397         var running int64
 398         for i := 0; i < runtime.NumCPU()*9/8+1; i++ {
 399                 tileJobs.Add(1)
 400                 atomic.AddInt64(&running, 1)
 401                 go func() {
 402                         defer tileJobs.Done()
 403                         defer atomic.AddInt64(&running, -1)
 404                         for fn := range todo {
 405                                 if len(errs) > 0 {
 406                                         return
 407                                 }
 408                                 err := fn()
 409                                 if err != nil {
 410                                         select {
 411                                         case errs <- err:
 412                                         default:
 413                                         }
 414                                 }
 415                                 remain := len(todo) + int(atomic.LoadInt64(&running)) - 1
 416                                 if remain < cap(todo) {
 417                                         ttl := time.Now().Sub(starttime) * time.Duration(remain) / time.Duration(cap(todo)-remain)
 418                                         eta := time.Now().Add(ttl)
 419                                         log.Printf("progress %d/%d, eta %v (%v)", cap(todo)-remain, cap(todo), eta, ttl)
 420                                 }
 421                         }
 422                 }()
 423         }
 424         tileJobs.Wait()
 425         encodeJobs.Wait()
 426
 427         go close(errs)
 428         err := <-errs
 429         if err != nil {
 430                 return err
 431         }
 432
 433         if cmd.outputStats != "" {
 434                 f, err := os.OpenFile(cmd.outputStats, os.O_CREATE|os.O_WRONLY, 0666)
 435                 if err != nil {
 436                         return err
 437                 }
 438                 var flatstats []importStats
 439                 for _, stats := range allstats {
 440                         flatstats = append(flatstats, stats...)
 441                 }
 442                 err = json.NewEncoder(f).Encode(flatstats)
 443                 if err != nil {
 444                         return err
 445                 }
 446         }
 447
 448         return nil
 449 }
 450
 451 func (cmd *importer) tileGVCF(tilelib *tileLibrary, infile string, phase int) (tileseq tileSeq, stats []importStats, err error) {
 452         if cmd.refFile == "" {
 453                 err = errors.New("cannot import vcf: reference data (-ref) not specified")
 454                 return
 455         }
 456         args := []string{"bcftools", "consensus", "--fasta-ref", cmd.refFile, "-H", fmt.Sprint(phase + 1), infile}
 457         indexsuffix := ".tbi"
 458         if _, err := os.Stat(infile + ".csi"); err == nil {
 459                 indexsuffix = ".csi"
 460         }
 461         if out, err := exec.Command("docker", "image", "ls", "-q", "lightning-runtime").Output(); err == nil && len(out) > 0 {
 462                 args = append([]string{
 463                         "docker", "run", "--rm",
 464                         "--log-driver=none",
 465                         "--volume=" + infile + ":" + infile + ":ro",
 466                         "--volume=" + infile + indexsuffix + ":" + infile + indexsuffix + ":ro",
 467                         "--volume=" + cmd.refFile + ":" + cmd.refFile + ":ro",
 468                         "lightning-runtime",
 469                 }, args...)
 470         }
 471         consensus := exec.Command(args[0], args[1:]...)
 472         consensus.Stderr = os.Stderr
 473         stdout, err := consensus.StdoutPipe()
 474         defer stdout.Close()
 475         if err != nil {
 476                 return
 477         }
 478         err = consensus.Start()
 479         if err != nil {
 480                 return
 481         }
 482         defer consensus.Wait()
 483         tileseq, stats, err = tilelib.TileFasta(fmt.Sprintf("%s phase %d", infile, phase+1), stdout)
 484         if err != nil {
 485                 return
 486         }
 487         err = stdout.Close()
 488         if err != nil {
 489                 return
 490         }
 491         err = consensus.Wait()
 492         if err != nil {
 493                 err = fmt.Errorf("%s phase %d: bcftools: %s", infile, phase, err)
 494                 return
 495         }
 496         return
 497 }
 498
 499 func flatten(variants [][]tileVariantID) []tileVariantID {
 500         ntags := 0
 501         for _, v := range variants {
 502                 if ntags < len(v) {
 503                         ntags = len(v)
 504                 }
 505         }
 506         flat := make([]tileVariantID, ntags*2)
 507         for i := 0; i < ntags; i++ {
 508                 for hap := 0; hap < 2; hap++ {
 509                         if i < len(variants[hap]) {
 510                                 flat[i*2+hap] = variants[hap][i]
 511                         }
 512                 }
 513         }
 514         return flat
 515 }
 516
 517 type importstatsplot struct{}
 518
 519 func (cmd *importstatsplot) RunCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
 520         err := cmd.Plot(stdin, stdout)
 521         if err != nil {
 522                 log.Errorf("%s", err)
 523                 return 1
 524         }
 525         return 0
 526 }
 527
 528 func (cmd *importstatsplot) Plot(stdin io.Reader, stdout io.Writer) error {
 529         var stats []importStats
 530         err := json.NewDecoder(stdin).Decode(&stats)
 531         if err != nil {
 532                 return err
 533         }
 534
 535         p, err := plot.New()
 536         if err != nil {
 537                 return err
 538         }
 539         p.Title.Text = "coverage preserved by import (excl X<0.65)"
 540         p.X.Label.Text = "input base calls ÷ sequence length"
 541         p.Y.Label.Text = "output base calls ÷ input base calls"
 542         p.Add(plotter.NewGrid())
 543
 544         data := map[string]plotter.XYs{}
 545         for _, stat := range stats {
 546                 data[stat.InputLabel] = append(data[stat.InputLabel], plotter.XY{
 547                         X: float64(stat.InputCoverage) / float64(stat.InputLength),
 548                         Y: float64(stat.TileCoverage) / float64(stat.InputCoverage),
 549                 })
 550         }
 551
 552         labels := []string{}
 553         for label := range data {
 554                 labels = append(labels, label)
 555         }
 556         sort.Strings(labels)
 557         palette, err := colorful.SoftPalette(len(labels))
 558         if err != nil {
 559                 return err
 560         }
 561         nextInPalette := 0
 562         for idx, label := range labels {
 563                 s, err := plotter.NewScatter(data[label])
 564                 if err != nil {
 565                         return err
 566                 }
 567                 s.GlyphStyle.Color = palette[idx]
 568                 s.GlyphStyle.Radius = vg.Millimeter / 2
 569                 s.GlyphStyle.Shape = draw.CrossGlyph{}
 570                 nextInPalette += 7
 571                 p.Add(s)
 572                 if false {
 573                         p.Legend.Add(label, s)
 574                 }
 575         }
 576         p.X.Min = 0.65
 577         p.X.Max = 1
 578
 579         w, err := p.WriterTo(8*vg.Inch, 6*vg.Inch, "svg")
 580         if err != nil {
 581                 return err
 582         }
 583         _, err = w.WriteTo(stdout)
 584         return err
 585 }