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         log "github.com/sirupsen/logrus"
  27 )
  28
  29 type importer struct {
  30         tagLibraryFile string
  31         refFile        string
  32         outputFile     string
  33         projectUUID    string
  34         runLocal       bool
  35         skipOOO        bool
  36         outputTiles    bool
  37         includeNoCalls bool
  38         outputStats    string
  39         encoder        *gob.Encoder
  40 }
  41
  42 func (cmd *importer) RunCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
  43         var err error
  44         defer func() {
  45                 if err != nil {
  46                         fmt.Fprintf(stderr, "%s\n", err)
  47                 }
  48         }()
  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 {
  65                 err = nil
  66                 return 0
  67         } else if err != nil {
  68                 return 2
  69         } else if cmd.tagLibraryFile == "" {
  70                 fmt.Fprintln(os.Stderr, "cannot import without -tag-library argument")
  71                 return 2
  72         } else if flags.NArg() == 0 {
  73                 flags.Usage()
  74                 return 2
  75         }
  76
  77         if *pprof != "" {
  78                 go func() {
  79                         log.Println(http.ListenAndServe(*pprof, nil))
  80                 }()
  81         }
  82
  83         lvl, err := log.ParseLevel(*loglevel)
  84         if err != nil {
  85                 return 2
  86         }
  87         log.SetLevel(lvl)
  88
  89         if !cmd.runLocal {
  90                 runner := arvadosContainerRunner{
  91                         Name:        "lightning import",
  92                         Client:      arvados.NewClientFromEnv(),
  93                         ProjectUUID: cmd.projectUUID,
  94                         RAM:         60000000000,
  95                         VCPUs:       16,
  96                         Priority:    *priority,
  97                 }
  98                 err = runner.TranslatePaths(&cmd.tagLibraryFile, &cmd.refFile, &cmd.outputFile)
  99                 if err != nil {
 100                         return 1
 101                 }
 102                 inputs := flags.Args()
 103                 for i := range inputs {
 104                         err = runner.TranslatePaths(&inputs[i])
 105                         if err != nil {
 106                                 return 1
 107                         }
 108                 }
 109                 if cmd.outputFile == "-" {
 110                         cmd.outputFile = "/mnt/output/library.gob"
 111                 } else {
 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")
 116                         return 1
 117                 }
 118                 runner.Args = append([]string{"import",
 119                         "-local=true",
 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,
 126                         "-ref", cmd.refFile,
 127                         "-o", cmd.outputFile,
 128                 }, inputs...)
 129                 var output string
 130                 output, err = runner.Run()
 131                 if err != nil {
 132                         return 1
 133                 }
 134                 fmt.Fprintln(stdout, output+"/library.gob")
 135                 return 0
 136         }
 137
 138         infiles, err := listInputFiles(flags.Args())
 139         if err != nil {
 140                 return 1
 141         }
 142
 143         taglib, err := cmd.loadTagLibrary()
 144         if err != nil {
 145                 return 1
 146         }
 147
 148         var output io.WriteCloser
 149         if cmd.outputFile == "-" {
 150                 output = nopCloser{stdout}
 151         } else {
 152                 output, err = os.OpenFile(cmd.outputFile, os.O_CREATE|os.O_WRONLY, 0777)
 153                 if err != nil {
 154                         return 1
 155                 }
 156                 defer output.Close()
 157         }
 158         bufw := bufio.NewWriter(output)
 159         cmd.encoder = gob.NewEncoder(bufw)
 160
 161         tilelib := &tileLibrary{taglib: taglib, includeNoCalls: cmd.includeNoCalls, skipOOO: cmd.skipOOO}
 162         if cmd.outputTiles {
 163                 cmd.encoder.Encode(LibraryEntry{TagSet: taglib.Tags()})
 164                 tilelib.encoder = cmd.encoder
 165         }
 166         go func() {
 167                 for range time.Tick(10 * time.Minute) {
 168                         log.Printf("tilelib.Len() == %d", tilelib.Len())
 169                 }
 170         }()
 171
 172         err = cmd.tileInputs(tilelib, infiles)
 173         if err != nil {
 174                 return 1
 175         }
 176         err = bufw.Flush()
 177         if err != nil {
 178                 return 1
 179         }
 180         err = output.Close()
 181         if err != nil {
 182                 return 1
 183         }
 184         return 0
 185 }
 186
 187 func (cmd *importer) tileFasta(tilelib *tileLibrary, infile string) (tileSeq, []importStats, error) {
 188         var input io.ReadCloser
 189         input, err := os.Open(infile)
 190         if err != nil {
 191                 return nil, nil, err
 192         }
 193         defer input.Close()
 194         if strings.HasSuffix(infile, ".gz") {
 195                 input, err = gzip.NewReader(input)
 196                 if err != nil {
 197                         return nil, nil, err
 198                 }
 199                 defer input.Close()
 200         }
 201         return tilelib.TileFasta(infile, input)
 202 }
 203
 204 func (cmd *importer) loadTagLibrary() (*tagLibrary, error) {
 205         log.Printf("tag library %s load starting", cmd.tagLibraryFile)
 206         f, err := os.Open(cmd.tagLibraryFile)
 207         if err != nil {
 208                 return nil, err
 209         }
 210         defer f.Close()
 211         var rdr io.ReadCloser = f
 212         if strings.HasSuffix(cmd.tagLibraryFile, ".gz") {
 213                 rdr, err = gzip.NewReader(f)
 214                 if err != nil {
 215                         return nil, fmt.Errorf("%s: gzip: %s", cmd.tagLibraryFile, err)
 216                 }
 217                 defer rdr.Close()
 218         }
 219         var taglib tagLibrary
 220         err = taglib.Load(rdr)
 221         if err != nil {
 222                 return nil, err
 223         }
 224         if taglib.Len() < 1 {
 225                 return nil, fmt.Errorf("cannot tile: tag library is empty")
 226         }
 227         log.Printf("tag library %s load done", cmd.tagLibraryFile)
 228         return &taglib, nil
 229 }
 230
 231 var (
 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)?$`)
 236 )
 237
 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)
 245                         }
 246                         continue
 247                 }
 248                 d, err := os.Open(path)
 249                 if err != nil {
 250                         return nil, fmt.Errorf("%s: open failed: %s", path, err)
 251                 }
 252                 defer d.Close()
 253                 names, err := d.Readdirnames(0)
 254                 if err != nil {
 255                         return nil, fmt.Errorf("%s: readdir failed: %s", path, err)
 256                 }
 257                 sort.Strings(names)
 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))
 263                         }
 264                 }
 265                 d.Close()
 266         }
 267         for _, file := range files {
 268                 if fastaFilenameRe.MatchString(file) {
 269                         continue
 270                 } else if vcfFilenameRe.MatchString(file) {
 271                         if _, err := os.Stat(file + ".csi"); err == nil {
 272                                 continue
 273                         } else if _, err = os.Stat(file + ".tbi"); err == nil {
 274                                 continue
 275                         } else {
 276                                 return nil, fmt.Errorf("%s: cannot read without .tbi or .csi index file", file)
 277                         }
 278                 } else {
 279                         return nil, fmt.Errorf("don't know how to handle filename %s", file)
 280                 }
 281         }
 282         return
 283 }
 284
 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
 294                 phases.Add(2)
 295                 variants := make([][]tileVariantID, 2)
 296                 if fasta1FilenameRe.MatchString(infile) {
 297                         todo <- func() error {
 298                                 defer phases.Done()
 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)
 306                                 return err
 307                         }
 308                         infile2 := fasta1FilenameRe.ReplaceAllString(infile, `.2.fa$1$2`)
 309                         todo <- func() error {
 310                                 defer phases.Done()
 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)
 318
 319                                 return err
 320                         }
 321                 } else if fastaFilenameRe.MatchString(infile) {
 322                         todo <- func() error {
 323                                 defer phases.Done()
 324                                 defer phases.Done()
 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
 329                                 if err != nil {
 330                                         return err
 331                                 }
 332                                 totlen := 0
 333                                 for _, tseq := range tseqs {
 334                                         totlen += len(tseq)
 335                                 }
 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}},
 339                                 })
 340                         }
 341                         // Don't write out a CompactGenomes entry
 342                         continue
 343                 } else if vcfFilenameRe.MatchString(infile) {
 344                         for phase := 0; phase < 2; phase++ {
 345                                 phase := phase
 346                                 todo <- func() error {
 347                                         defer phases.Done()
 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)
 355                                         return err
 356                                 }
 357                         }
 358                 } else {
 359                         panic(fmt.Sprintf("bug: unhandled filename %q", infile))
 360                 }
 361                 encodeJobs.Add(1)
 362                 go func() {
 363                         defer encodeJobs.Done()
 364                         phases.Wait()
 365                         if len(errs) > 0 {
 366                                 return
 367                         }
 368                         err := cmd.encoder.Encode(LibraryEntry{
 369                                 CompactGenomes: []CompactGenome{{Name: infile, Variants: flatten(variants)}},
 370                         })
 371                         if err != nil {
 372                                 select {
 373                                 case errs <- err:
 374                                 default:
 375                                 }
 376                         }
 377                 }()
 378         }
 379         go close(todo)
 380         var tileJobs sync.WaitGroup
 381         var running int64
 382         for i := 0; i < runtime.NumCPU()*9/8+1; i++ {
 383                 tileJobs.Add(1)
 384                 atomic.AddInt64(&running, 1)
 385                 go func() {
 386                         defer tileJobs.Done()
 387                         defer atomic.AddInt64(&running, -1)
 388                         for fn := range todo {
 389                                 if len(errs) > 0 {
 390                                         return
 391                                 }
 392                                 err := fn()
 393                                 if err != nil {
 394                                         select {
 395                                         case errs <- err:
 396                                         default:
 397                                         }
 398                                 }
 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)
 404                                 }
 405                         }
 406                 }()
 407         }
 408         tileJobs.Wait()
 409         encodeJobs.Wait()
 410
 411         go close(errs)
 412         err := <-errs
 413         if err != nil {
 414                 return err
 415         }
 416
 417         if cmd.outputStats != "" {
 418                 f, err := os.OpenFile(cmd.outputStats, os.O_CREATE|os.O_WRONLY, 0666)
 419                 if err != nil {
 420                         return err
 421                 }
 422                 var flatstats []importStats
 423                 for _, stats := range allstats {
 424                         flatstats = append(flatstats, stats...)
 425                 }
 426                 err = json.NewEncoder(f).Encode(flatstats)
 427                 if err != nil {
 428                         return err
 429                 }
 430         }
 431
 432         return nil
 433 }
 434
 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")
 438                 return
 439         }
 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 {
 443                 indexsuffix = ".csi"
 444         }
 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",
 448                         "--log-driver=none",
 449                         "--volume=" + infile + ":" + infile + ":ro",
 450                         "--volume=" + infile + indexsuffix + ":" + infile + indexsuffix + ":ro",
 451                         "--volume=" + cmd.refFile + ":" + cmd.refFile + ":ro",
 452                         "lightning-runtime",
 453                 }, args...)
 454         }
 455         consensus := exec.Command(args[0], args[1:]...)
 456         consensus.Stderr = os.Stderr
 457         stdout, err := consensus.StdoutPipe()
 458         defer stdout.Close()
 459         if err != nil {
 460                 return
 461         }
 462         err = consensus.Start()
 463         if err != nil {
 464                 return
 465         }
 466         defer consensus.Wait()
 467         tileseq, stats, err = tilelib.TileFasta(fmt.Sprintf("%s phase %d", infile, phase+1), stdout)
 468         if err != nil {
 469                 return
 470         }
 471         err = stdout.Close()
 472         if err != nil {
 473                 return
 474         }
 475         err = consensus.Wait()
 476         if err != nil {
 477                 err = fmt.Errorf("%s phase %d: bcftools: %s", infile, phase, err)
 478                 return
 479         }
 480         return
 481 }
 482
 483 func flatten(variants [][]tileVariantID) []tileVariantID {
 484         ntags := 0
 485         for _, v := range variants {
 486                 if ntags < len(v) {
 487                         ntags = len(v)
 488                 }
 489         }
 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]
 495                         }
 496                 }
 497         }
 498         return flat
 499 }