import.go

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