vcf2fasta.go

   1 package main
   2
   3 import (
   4         "bufio"
   5         "bytes"
   6         "compress/gzip"
   7         "context"
   8         "errors"
   9         "flag"
  10         "fmt"
  11         "io"
  12         "io/ioutil"
  13         "net/http"
  14         _ "net/http/pprof"
  15         "os"
  16         "os/exec"
  17         "path/filepath"
  18         "regexp"
  19         "runtime"
  20         "strconv"
  21         "strings"
  22         "sync"
  23         "syscall"
  24
  25         "github.com/klauspost/pgzip"
  26         log "github.com/sirupsen/logrus"
  27 )
  28
  29 type vcf2fasta struct {
  30         refFile           string
  31         genomeFile        string
  32         mask              bool
  33         gvcfRegionsPy     string
  34         gvcfRegionsPyData []byte
  35         gvcfType          string
  36         projectUUID       string
  37         outputDir         string
  38         runLocal          bool
  39         vcpus             int
  40         batchArgs
  41
  42         stderr io.Writer
  43 }
  44
  45 func (cmd *vcf2fasta) RunCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
  46         var err error
  47         defer func() {
  48                 if err != nil {
  49                         fmt.Fprintf(stderr, "%s\n", err)
  50                 }
  51         }()
  52         flags := flag.NewFlagSet("", flag.ContinueOnError)
  53         flags.SetOutput(stderr)
  54         flags.StringVar(&cmd.refFile, "ref", "", "reference fasta `file`")
  55         flags.StringVar(&cmd.genomeFile, "genome", "", "reference genome `file`")
  56         flags.BoolVar(&cmd.mask, "mask", false, "mask uncalled regions (default: output hom ref)")
  57         flags.StringVar(&cmd.gvcfRegionsPy, "gvcf-regions.py", "https://raw.githubusercontent.com/lijiayong/gvcf_regions/master/gvcf_regions.py", "source of gvcf_regions.py")
  58         flags.StringVar(&cmd.gvcfType, "gvcf-type", "gatk", "gvcf_type argument to gvcf_regions.py: gatk, complete_genomics, freebayes")
  59         flags.StringVar(&cmd.projectUUID, "project", "", "project `UUID` for containers and output data")
  60         flags.StringVar(&cmd.outputDir, "output-dir", "", "output directory")
  61         flags.IntVar(&cmd.vcpus, "vcpus", 0, "number of VCPUs to request for arvados container (default: 2*number of input files, max 32)")
  62         flags.BoolVar(&cmd.runLocal, "local", false, "run on local host (default: run in an arvados container)")
  63         cmd.batchArgs.Flags(flags)
  64         priority := flags.Int("priority", 500, "container request priority")
  65         pprof := flags.String("pprof", "", "serve Go profile data at http://`[addr]:port`")
  66         err = flags.Parse(args)
  67         if err == flag.ErrHelp {
  68                 err = nil
  69                 return 0
  70         } else if err != nil {
  71                 return 2
  72         } else if cmd.refFile == "" {
  73                 err = errors.New("reference data (-ref) not specified")
  74                 return 2
  75         } else if flags.NArg() == 0 {
  76                 flags.Usage()
  77                 return 2
  78         }
  79         cmd.stderr = stderr
  80
  81         if *pprof != "" {
  82                 go func() {
  83                         log.Println(http.ListenAndServe(*pprof, nil))
  84                 }()
  85         }
  86
  87         if cmd.mask {
  88                 err = cmd.loadRegionsPy()
  89                 if err != nil {
  90                         return 1
  91                 }
  92         }
  93
  94         if !cmd.runLocal {
  95                 if cmd.outputDir != "" {
  96                         err = errors.New("cannot specify output dir in non-local mode")
  97                         return 2
  98                 }
  99                 if cmd.vcpus < 1 {
 100                         var infiles []string
 101                         infiles, err = listInputFiles(flags.Args())
 102                         if err != nil {
 103                                 return 1
 104                         }
 105                         batchsize := (len(infiles) + cmd.batchArgs.batches - 1) / cmd.batchArgs.batches
 106                         if cmd.vcpus = batchsize * 2; cmd.vcpus > 32 {
 107                                 cmd.vcpus = 32
 108                         }
 109                 }
 110                 runner := arvadosContainerRunner{
 111                         Name:        "lightning vcf2fasta",
 112                         Client:      arvadosClientFromEnv,
 113                         ProjectUUID: cmd.projectUUID,
 114                         RAM:         2<<30 + int64(cmd.vcpus)<<28,
 115                         VCPUs:       cmd.vcpus,
 116                         Priority:    *priority,
 117                         KeepCache:   2,
 118                         APIAccess:   true,
 119                         Mounts: map[string]map[string]interface{}{
 120                                 "/gvcf_regions.py": map[string]interface{}{
 121                                         "kind":    "text",
 122                                         "content": string(cmd.gvcfRegionsPyData),
 123                                 },
 124                         },
 125                 }
 126                 err = runner.TranslatePaths(&cmd.refFile, &cmd.genomeFile)
 127                 if err != nil {
 128                         return 1
 129                 }
 130                 inputs := flags.Args()
 131                 for i := range inputs {
 132                         err = runner.TranslatePaths(&inputs[i])
 133                         if err != nil {
 134                                 return 1
 135                         }
 136                 }
 137                 var outputs []string
 138                 outputs, err = cmd.batchArgs.RunBatches(context.Background(), func(ctx context.Context, batch int) (string, error) {
 139                         runner := runner
 140                         if cmd.mask {
 141                                 runner.RAM += int64(cmd.vcpus) << 31
 142                         }
 143                         runner.Args = []string{"vcf2fasta",
 144                                 "-local=true",
 145                                 "-ref", cmd.refFile, fmt.Sprintf("-mask=%v", cmd.mask),
 146                                 "-genome", cmd.genomeFile,
 147                                 "-gvcf-regions.py", "/gvcf_regions.py",
 148                                 "-gvcf-type", cmd.gvcfType,
 149                                 "-output-dir", "/mnt/output",
 150                         }
 151                         runner.Args = append(runner.Args, cmd.batchArgs.Args(batch)...)
 152                         runner.Args = append(runner.Args, inputs...)
 153                         log.Printf("batch %d: %v", batch, runner.Args)
 154                         return runner.RunContext(ctx)
 155                 })
 156                 if err != nil {
 157                         return 1
 158                 }
 159                 fmt.Fprintln(stdout, strings.Join(outputs, " "))
 160                 return 0
 161         }
 162
 163         infiles, err := listInputFiles(flags.Args())
 164         if err != nil {
 165                 return 1
 166         }
 167         infiles = cmd.batchArgs.Slice(infiles)
 168
 169         type job struct {
 170                 vcffile string
 171                 phase   int
 172         }
 173         todo := make(chan job)
 174         go func() {
 175                 for _, infile := range infiles {
 176                         for phase := 1; phase <= 2; phase++ {
 177                                 todo <- job{vcffile: infile, phase: phase}
 178                         }
 179                 }
 180                 close(todo)
 181         }()
 182
 183         done := make(chan error, runtime.NumCPU()*2)
 184         var wg sync.WaitGroup
 185         for i := 0; i < runtime.NumCPU(); i++ {
 186                 wg.Add(1)
 187                 go func() {
 188                         defer wg.Done()
 189                         for job := range todo {
 190                                 if len(done) > 0 {
 191                                         // a different worker encountered an error
 192                                         return
 193                                 }
 194                                 err := cmd.vcf2fasta(job.vcffile, job.phase)
 195                                 if err != nil {
 196                                         done <- fmt.Errorf("%s phase %d: %s", job.vcffile, job.phase, err)
 197                                         return
 198                                 }
 199                         }
 200                 }()
 201         }
 202         go func() {
 203                 wg.Wait()
 204                 close(done)
 205         }()
 206
 207         err = <-done
 208         if err != nil {
 209                 return 1
 210         }
 211         return 0
 212 }
 213
 214 func maybeInDocker(args, mountfiles []string) []string {
 215         if out, err := exec.Command("docker", "image", "ls", "-q", "lightning-runtime").Output(); err != nil || len(out) == 0 {
 216                 return args
 217         }
 218         dockerrun := []string{
 219                 "docker", "run", "--rm", "-i",
 220                 "--log-driver=none",
 221         }
 222         for _, f := range mountfiles {
 223                 dockerrun = append(dockerrun, "--volume="+f+":"+f+":ro")
 224         }
 225         dockerrun = append(dockerrun, "lightning-runtime")
 226         dockerrun = append(dockerrun, args...)
 227         return dockerrun
 228 }
 229
 230 func (cmd *vcf2fasta) vcf2fasta(infile string, phase int) error {
 231         ctx, cancel := context.WithCancel(context.Background())
 232         defer cancel()
 233
 234         _, basename := filepath.Split(infile)
 235         outfile := filepath.Join(cmd.outputDir, fmt.Sprintf("%s.%d.fasta.gz", basename, phase))
 236         outf, err := os.OpenFile(outfile, os.O_CREATE|os.O_WRONLY, 0777)
 237         if err != nil {
 238                 return fmt.Errorf("error opening output file: %s", err)
 239         }
 240         defer outf.Close()
 241         bufw := bufio.NewWriterSize(outf, 8*1024*1024)
 242         gzipw := pgzip.NewWriter(bufw)
 243         defer gzipw.Close()
 244
 245         var maskfifo string // filename of mask fifo if we're running bedtools, otherwise ""
 246
 247         var wg sync.WaitGroup
 248         errs := make(chan error, 2)
 249         if cmd.mask {
 250                 chrSize := map[string]int{}
 251
 252                 vcffile, err := open(infile)
 253                 if err != nil {
 254                         return err
 255                 }
 256                 defer vcffile.Close()
 257                 var rdr io.Reader = vcffile
 258                 rdr = bufio.NewReaderSize(rdr, 8*1024*1024)
 259                 if strings.HasSuffix(infile, ".gz") {
 260                         rdr, err = gzip.NewReader(vcffile)
 261                         if err != nil {
 262                                 return err
 263                         }
 264                 }
 265                 contigre := regexp.MustCompile(`([^=,]*)=([^>,]*)`)
 266                 scanner := bufio.NewScanner(rdr)
 267                 for scanner.Scan() {
 268                         if s := scanner.Text(); !strings.HasPrefix(s, "##") {
 269                                 break
 270                         } else if !strings.HasPrefix(s, "##contig=<") {
 271                                 continue
 272                         } else {
 273                                 kv := map[string]string{}
 274                                 for _, m := range contigre.FindAllStringSubmatch(s[10:], -1) {
 275                                         kv[m[1]] = m[2]
 276                                 }
 277                                 if kv["ID"] != "" && kv["length"] != "" {
 278                                         chrSize[kv["ID"]], _ = strconv.Atoi(kv["length"])
 279                                 }
 280                         }
 281                 }
 282                 if err = scanner.Err(); err != nil {
 283                         return fmt.Errorf("error scanning input file %q: %s", infile, err)
 284                 }
 285
 286                 var regions bytes.Buffer
 287                 bedargs := []string{"python2", "-"}
 288                 if cmd.gvcfType != "" {
 289                         bedargs = append(bedargs, "--gvcf_type", cmd.gvcfType)
 290                 }
 291                 bedargs = append(bedargs, infile)
 292                 bed := exec.CommandContext(ctx, bedargs[0], bedargs[1:]...)
 293                 bed.Stdin = bytes.NewBuffer(cmd.gvcfRegionsPyData)
 294                 bed.Stdout = &regions
 295                 bed.Stderr = cmd.stderr
 296                 log.Printf("running %v", bed.Args)
 297                 err = bed.Run()
 298                 log.Printf("exited %v", bed.Args)
 299                 if err != nil {
 300                         return fmt.Errorf("gvcf_regions: %s", err)
 301                 }
 302
 303                 if cmd.genomeFile != "" {
 304                         // Read chromosome sizes from genome file in
 305                         // case any weren't specified in the VCF
 306                         // header.
 307                         genomeFile, err := open(cmd.genomeFile)
 308                         if err != nil {
 309                                 return fmt.Errorf("error opening genome file %q: %s", cmd.genomeFile, err)
 310                         }
 311                         scanner := bufio.NewScanner(genomeFile)
 312                         for scanner.Scan() {
 313                                 var chr string
 314                                 var size int
 315                                 _, err := fmt.Sscanf(scanner.Text(), "%s\t%d", &chr, &size)
 316                                 if err != nil {
 317                                         return fmt.Errorf("error parsing genome file %q: %s", cmd.genomeFile, err)
 318                                 }
 319                                 if chrSize[chr] == 0 {
 320                                         chrSize[chr] = size
 321                                 }
 322                         }
 323                         if err = scanner.Err(); err != nil {
 324                                 return fmt.Errorf("error scanning genome file %q: %s", cmd.genomeFile, err)
 325                         }
 326                 }
 327
 328                 // "bedtools complement" expects the chromosome sizes
 329                 // ("genome file") to appear in the same order as the
 330                 // chromosomes in the input vcf, so we need to sort
 331                 // them.
 332                 scanner = bufio.NewScanner(bytes.NewBuffer(append([]byte(nil), regions.Bytes()...)))
 333                 var sortedGenomeFile bytes.Buffer
 334                 for scanner.Scan() {
 335                         var chr string
 336                         var size int
 337                         _, err := fmt.Sscanf(scanner.Text(), "%s\t%d", &chr, &size)
 338                         if err != nil {
 339                                 return fmt.Errorf("error parsing gvcf_regions output: %s", err)
 340                         }
 341                         if size, ok := chrSize[chr]; ok {
 342                                 fmt.Fprintf(&sortedGenomeFile, "%s\t%d\n", chr, size)
 343                                 delete(chrSize, chr)
 344                         }
 345                 }
 346
 347                 // The bcftools --mask argument needs to end in ".bed"
 348                 // in order to be parsed as a BED file, so we need to
 349                 // use a named pipe instead of stdin.
 350                 tempdir, err := ioutil.TempDir("", "")
 351                 if err != nil {
 352                         return fmt.Errorf("TempDir: %s", err)
 353                 }
 354                 defer os.RemoveAll(tempdir)
 355                 maskfifo = filepath.Join(tempdir, "fifo.bed")
 356                 err = syscall.Mkfifo(maskfifo, 0600)
 357                 if err != nil {
 358                         return fmt.Errorf("mkfifo: %s", err)
 359                 }
 360
 361                 // bedtools complement can't seem to read from a pipe
 362                 // reliably -- "Error: line number 1 of file
 363                 // /dev/stdin has 1 fields, but 3 were expected." --
 364                 // so we stage to a temp file.
 365                 regionsFile := filepath.Join(tempdir, "gvcf_regions.bed")
 366                 err = ioutil.WriteFile(regionsFile, regions.Bytes(), 0644)
 367                 if err != nil {
 368                         return err
 369                 }
 370
 371                 wg.Add(1)
 372                 go func() {
 373                         defer wg.Done()
 374
 375                         maskfifow, err := os.OpenFile(maskfifo, os.O_WRONLY, 0)
 376                         if err != nil {
 377                                 errs <- err
 378                                 return
 379                         }
 380                         defer maskfifow.Close()
 381
 382                         bedcompargs := []string{"bedtools", "complement", "-i", regionsFile, "-g", "/dev/stdin"}
 383                         bedcompargs = maybeInDocker(bedcompargs, []string{cmd.genomeFile})
 384                         bedcomp := exec.CommandContext(ctx, bedcompargs[0], bedcompargs[1:]...)
 385                         bedcomp.Stdin = &sortedGenomeFile
 386                         bedcomp.Stdout = maskfifow
 387                         bedcomp.Stderr = cmd.stderr
 388                         log.Printf("running %v", bedcomp.Args)
 389                         err = bedcomp.Run()
 390                         log.Printf("exited %v", bedcomp.Args)
 391                         if err != nil {
 392                                 errs <- fmt.Errorf("bedtools complement: %s", err)
 393                                 return
 394                         }
 395                         err = maskfifow.Close()
 396                         if err != nil {
 397                                 errs <- err
 398                                 return
 399                         }
 400                 }()
 401         }
 402
 403         wg.Add(1)
 404         go func() {
 405                 defer wg.Done()
 406                 consargs := []string{"bcftools", "consensus", "--fasta-ref", cmd.refFile, "-H", fmt.Sprint(phase)}
 407                 if maskfifo != "" {
 408                         consargs = append(consargs, "--mask", maskfifo)
 409                 }
 410                 consargs = append(consargs, infile)
 411                 indexsuffix := ".tbi"
 412                 if _, err := os.Stat(infile + ".csi"); err == nil {
 413                         indexsuffix = ".csi"
 414                 }
 415                 mounts := []string{infile, infile + indexsuffix, cmd.refFile}
 416                 if maskfifo != "" {
 417                         mounts = append(mounts, maskfifo)
 418                 }
 419                 consargs = maybeInDocker(consargs, mounts)
 420
 421                 consensus := exec.CommandContext(ctx, consargs[0], consargs[1:]...)
 422                 consensus.Stderr = os.Stderr
 423                 consensus.Stdout = gzipw
 424                 consensus.Stderr = cmd.stderr
 425                 log.Printf("running %v", consensus.Args)
 426                 err = consensus.Run()
 427                 if err != nil {
 428                         errs <- fmt.Errorf("bcftools consensus: %s", err)
 429                         return
 430                 }
 431                 log.Printf("exited %v", consensus.Args)
 432                 err = gzipw.Close()
 433                 if err != nil {
 434                         errs <- err
 435                         return
 436                 }
 437                 err = bufw.Flush()
 438                 if err != nil {
 439                         errs <- err
 440                         return
 441                 }
 442                 errs <- outf.Close()
 443         }()
 444
 445         go func() {
 446                 wg.Wait()
 447                 close(errs)
 448         }()
 449
 450         for err := range errs {
 451                 if err != nil {
 452                         cancel()
 453                         wg.Wait()
 454                         return err
 455                 }
 456         }
 457         return nil
 458 }
 459
 460 func (cmd *vcf2fasta) loadRegionsPy() error {
 461         if strings.HasPrefix(cmd.gvcfRegionsPy, "http") {
 462                 resp, err := http.Get(cmd.gvcfRegionsPy)
 463                 if err != nil {
 464                         return err
 465                 }
 466                 if resp.StatusCode != http.StatusOK {
 467                         return fmt.Errorf("get %q: http status %d", cmd.gvcfRegionsPy, resp.StatusCode)
 468                 }
 469                 buf, err := ioutil.ReadAll(resp.Body)
 470                 if err != nil {
 471                         return fmt.Errorf("get %q: read body: %s", cmd.gvcfRegionsPy, err)
 472                 }
 473                 cmd.gvcfRegionsPyData = buf
 474                 return nil
 475         } else {
 476                 buf, err := ioutil.ReadFile(cmd.gvcfRegionsPy)
 477                 if err != nil {
 478                         return err
 479                 }
 480                 cmd.gvcfRegionsPyData = buf
 481                 return nil
 482         }
 483 }