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         "log"
  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 )
  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                 err = runner.Run()
 102                 if err != nil {
 103                         return 1
 104                 }
 105                 return 0
 106         }
 107
 108         infiles, err := listInputFiles(flags.Args())
 109         if err != nil {
 110                 return 1
 111         }
 112
 113         tilelib, err := cmd.loadTileLibrary()
 114         if err != nil {
 115                 return 1
 116         }
 117         go func() {
 118                 for range time.Tick(10 * time.Second) {
 119                         log.Printf("tilelib.Len() == %d", tilelib.Len())
 120                 }
 121         }()
 122
 123         var output io.WriteCloser
 124         if cmd.outputFile == "" {
 125                 output = nopCloser{stdout}
 126         } else {
 127                 output, err = os.OpenFile(cmd.outputFile, os.O_CREATE|os.O_WRONLY, 0777)
 128                 if err != nil {
 129                         return 1
 130                 }
 131                 defer output.Close()
 132         }
 133         bufw := bufio.NewWriter(output)
 134         cmd.encoder = gob.NewEncoder(bufw)
 135
 136         err = cmd.tileInputs(tilelib, infiles)
 137         if err != nil {
 138                 return 1
 139         }
 140         err = bufw.Flush()
 141         if err != nil {
 142                 return 1
 143         }
 144         err = output.Close()
 145         if err != nil {
 146                 return 1
 147         }
 148         return 0
 149 }
 150
 151 func (cmd *importer) tileFasta(tilelib *tileLibrary, infile string) (tileSeq, error) {
 152         var input io.ReadCloser
 153         input, err := os.Open(infile)
 154         if err != nil {
 155                 return nil, err
 156         }
 157         defer input.Close()
 158         if strings.HasSuffix(infile, ".gz") {
 159                 input, err = gzip.NewReader(input)
 160                 if err != nil {
 161                         return nil, err
 162                 }
 163                 defer input.Close()
 164         }
 165         return tilelib.TileFasta(infile, input)
 166 }
 167
 168 func (cmd *importer) loadTileLibrary() (*tileLibrary, error) {
 169         log.Printf("tag library %s load starting", cmd.tagLibraryFile)
 170         f, err := os.Open(cmd.tagLibraryFile)
 171         if err != nil {
 172                 return nil, err
 173         }
 174         defer f.Close()
 175         var rdr io.ReadCloser = f
 176         if strings.HasSuffix(cmd.tagLibraryFile, ".gz") {
 177                 rdr, err = gzip.NewReader(f)
 178                 if err != nil {
 179                         return nil, fmt.Errorf("%s: gzip: %s", cmd.tagLibraryFile, err)
 180                 }
 181                 defer rdr.Close()
 182         }
 183         var taglib tagLibrary
 184         err = taglib.Load(rdr)
 185         if err != nil {
 186                 return nil, err
 187         }
 188         if taglib.Len() < 1 {
 189                 return nil, fmt.Errorf("cannot tile: tag library is empty")
 190         }
 191         log.Printf("tag library %s load done", cmd.tagLibraryFile)
 192         return &tileLibrary{taglib: &taglib}, nil
 193 }
 194
 195 func listInputFiles(paths []string) (files []string, err error) {
 196         for _, path := range paths {
 197                 if fi, err := os.Stat(path); err != nil {
 198                         return nil, fmt.Errorf("%s: stat failed: %s", path, err)
 199                 } else if !fi.IsDir() {
 200                         if !strings.HasSuffix(path, ".2.fasta") || strings.HasSuffix(path, ".2.fasta.gz") {
 201                                 files = append(files, path)
 202                         }
 203                         continue
 204                 }
 205                 d, err := os.Open(path)
 206                 if err != nil {
 207                         return nil, fmt.Errorf("%s: open failed: %s", path, err)
 208                 }
 209                 defer d.Close()
 210                 names, err := d.Readdirnames(0)
 211                 if err != nil {
 212                         return nil, fmt.Errorf("%s: readdir failed: %s", path, err)
 213                 }
 214                 sort.Strings(names)
 215                 for _, name := range names {
 216                         if strings.HasSuffix(name, ".vcf") || strings.HasSuffix(name, ".vcf.gz") {
 217                                 files = append(files, filepath.Join(path, name))
 218                         } else if strings.HasSuffix(name, ".1.fasta") || strings.HasSuffix(name, ".1.fasta.gz") {
 219                                 files = append(files, filepath.Join(path, name))
 220                         }
 221                 }
 222                 d.Close()
 223         }
 224         for _, file := range files {
 225                 if strings.HasSuffix(file, ".1.fasta") || strings.HasSuffix(file, ".1.fasta.gz") {
 226                         continue
 227                 } else if _, err := os.Stat(file + ".csi"); err == nil {
 228                         continue
 229                 } else if _, err = os.Stat(file + ".tbi"); err == nil {
 230                         continue
 231                 } else {
 232                         return nil, fmt.Errorf("%s: cannot read without .tbi or .csi index file", file)
 233                 }
 234         }
 235         return
 236 }
 237
 238 func (cmd *importer) tileInputs(tilelib *tileLibrary, infiles []string) error {
 239         starttime := time.Now()
 240         errs := make(chan error, 1)
 241         todo := make(chan func() error, len(infiles)*2)
 242         var encodeJobs sync.WaitGroup
 243         for _, infile := range infiles {
 244                 infile := infile
 245                 var phases sync.WaitGroup
 246                 phases.Add(2)
 247                 variants := make([][]tileVariantID, 2)
 248                 if strings.HasSuffix(infile, ".1.fasta") || strings.HasSuffix(infile, ".1.fasta.gz") {
 249                         todo <- func() error {
 250                                 defer phases.Done()
 251                                 log.Printf("%s starting", infile)
 252                                 defer log.Printf("%s done", infile)
 253                                 tseqs, err := cmd.tileFasta(tilelib, infile)
 254                                 variants[0] = tseqs.Variants()
 255                                 return err
 256                         }
 257                         infile2 := regexp.MustCompile(`\.1\.fasta(\.gz)?$`).ReplaceAllString(infile, `.2.fasta$1`)
 258                         todo <- func() error {
 259                                 defer phases.Done()
 260                                 log.Printf("%s starting", infile2)
 261                                 defer log.Printf("%s done", infile2)
 262                                 tseqs, err := cmd.tileFasta(tilelib, infile2)
 263                                 variants[1] = tseqs.Variants()
 264                                 return err
 265                         }
 266                 } else {
 267                         for phase := 0; phase < 2; phase++ {
 268                                 phase := phase
 269                                 todo <- func() error {
 270                                         defer phases.Done()
 271                                         log.Printf("%s phase %d starting", infile, phase+1)
 272                                         defer log.Printf("%s phase %d done", infile, phase+1)
 273                                         tseqs, err := cmd.tileGVCF(tilelib, infile, phase)
 274                                         variants[phase] = tseqs.Variants()
 275                                         return err
 276                                 }
 277                         }
 278                 }
 279                 encodeJobs.Add(1)
 280                 go func() {
 281                         defer encodeJobs.Done()
 282                         phases.Wait()
 283                         if len(errs) > 0 {
 284                                 return
 285                         }
 286                         ntags := len(variants[0])
 287                         if ntags < len(variants[1]) {
 288                                 ntags = len(variants[1])
 289                         }
 290                         flat := make([]tileVariantID, ntags*2)
 291                         for i := 0; i < ntags; i++ {
 292                                 flat[i*2] = variants[0][i]
 293                                 flat[i*2+1] = variants[1][i]
 294                         }
 295                         err := cmd.encoder.Encode(LibraryEntry{
 296                                 CompactGenomes: []CompactGenome{{Name: infile, Variants: flat}},
 297                         })
 298                         if err != nil {
 299                                 select {
 300                                 case errs <- err:
 301                                 default:
 302                                 }
 303                         }
 304                 }()
 305         }
 306         go close(todo)
 307         var tileJobs sync.WaitGroup
 308         running := int64(runtime.NumCPU())
 309         for i := 0; i < runtime.NumCPU(); i++ {
 310                 tileJobs.Add(1)
 311                 go func() {
 312                         defer tileJobs.Done()
 313                         defer atomic.AddInt64(&running, -1)
 314                         for fn := range todo {
 315                                 if len(errs) > 0 {
 316                                         return
 317                                 }
 318                                 err := fn()
 319                                 if err != nil {
 320                                         select {
 321                                         case errs <- err:
 322                                         default:
 323                                         }
 324                                 }
 325                                 remain := len(todo) + int(atomic.LoadInt64(&running)) - 1
 326                                 ttl := time.Now().Sub(starttime) * time.Duration(remain) / time.Duration(cap(todo)-remain)
 327                                 eta := time.Now().Add(ttl)
 328                                 log.Printf("progress %d/%d, eta %v (%v)", cap(todo)-remain, cap(todo), eta, ttl)
 329                         }
 330                 }()
 331         }
 332         tileJobs.Wait()
 333         encodeJobs.Wait()
 334         go close(errs)
 335         return <-errs
 336 }
 337
 338 func (cmd *importer) tileGVCF(tilelib *tileLibrary, infile string, phase int) (tileseq tileSeq, err error) {
 339         if cmd.refFile == "" {
 340                 err = errors.New("cannot import vcf: reference data (-ref) not specified")
 341                 return
 342         }
 343         args := []string{"bcftools", "consensus", "--fasta-ref", cmd.refFile, "-H", fmt.Sprint(phase + 1), infile}
 344         indexsuffix := ".tbi"
 345         if _, err := os.Stat(infile + ".csi"); err == nil {
 346                 indexsuffix = ".csi"
 347         }
 348         if out, err := exec.Command("docker", "image", "ls", "-q", "lightning-runtime").Output(); err == nil && len(out) > 0 {
 349                 args = append([]string{
 350                         "docker", "run", "--rm",
 351                         "--log-driver=none",
 352                         "--volume=" + infile + ":" + infile + ":ro",
 353                         "--volume=" + infile + indexsuffix + ":" + infile + indexsuffix + ":ro",
 354                         "--volume=" + cmd.refFile + ":" + cmd.refFile + ":ro",
 355                         "lightning-runtime",
 356                 }, args...)
 357         }
 358         consensus := exec.Command(args[0], args[1:]...)
 359         consensus.Stderr = os.Stderr
 360         stdout, err := consensus.StdoutPipe()
 361         defer stdout.Close()
 362         if err != nil {
 363                 return
 364         }
 365         err = consensus.Start()
 366         if err != nil {
 367                 return
 368         }
 369         defer consensus.Wait()
 370         tileseq, err = tilelib.TileFasta(fmt.Sprintf("%s phase %d", infile, phase+1), stdout)
 371         if err != nil {
 372                 return
 373         }
 374         err = stdout.Close()
 375         if err != nil {
 376                 return
 377         }
 378         err = consensus.Wait()
 379         if err != nil {
 380                 err = fmt.Errorf("%s phase %d: bcftools: %s", infile, phase, err)
 381                 return
 382         }
 383         return
 384 }