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