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