19527: Load training-set flag from samples.csv.
[lightning.git] / dumpgob.go
1 // Copyright (C) The Lightning Authors. All rights reserved.
2 //
3 // SPDX-License-Identifier: AGPL-3.0
4
5 package lightning
6
7 import (
8         "bufio"
9         "errors"
10         "flag"
11         "fmt"
12         "io"
13         "net/http"
14         _ "net/http/pprof"
15         "os"
16         "strings"
17
18         "git.arvados.org/arvados.git/sdk/go/arvados"
19         log "github.com/sirupsen/logrus"
20 )
21
22 type dumpGob struct{}
23
24 func (cmd *dumpGob) RunCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
25         var err error
26         defer func() {
27                 if err != nil {
28                         fmt.Fprintf(stderr, "%s\n", err)
29                 }
30         }()
31         flags := flag.NewFlagSet("", flag.ContinueOnError)
32         flags.SetOutput(stderr)
33         pprof := flags.String("pprof", "", "serve Go profile data at http://`[addr]:port`")
34         runlocal := flags.Bool("local", false, "run on local host (default: run in an arvados container)")
35         projectUUID := flags.String("project", "", "project `UUID` for output data")
36         priority := flags.Int("priority", 500, "container request priority")
37         inputFilename := flags.String("i", "-", "input `file` (library)")
38         outputFilename := flags.String("o", "-", "output `file`")
39         err = flags.Parse(args)
40         if err == flag.ErrHelp {
41                 err = nil
42                 return 0
43         } else if err != nil {
44                 return 2
45         }
46
47         if *pprof != "" {
48                 go func() {
49                         log.Println(http.ListenAndServe(*pprof, nil))
50                 }()
51         }
52
53         if !*runlocal {
54                 if *outputFilename != "-" {
55                         err = errors.New("cannot specify output file in container mode: not implemented")
56                         return 1
57                 }
58                 runner := arvadosContainerRunner{
59                         Name:        "lightning dumpgob",
60                         Client:      arvados.NewClientFromEnv(),
61                         ProjectUUID: *projectUUID,
62                         RAM:         4000000000,
63                         VCPUs:       1,
64                         Priority:    *priority,
65                 }
66                 err = runner.TranslatePaths(inputFilename)
67                 if err != nil {
68                         return 1
69                 }
70                 runner.Args = []string{"dumpgob", "-local=true", fmt.Sprintf("-pprof=%v", *pprof), "-i", *inputFilename, "-o", "/mnt/output/dumpgob.txt"}
71                 var output string
72                 output, err = runner.Run()
73                 if err != nil {
74                         return 1
75                 }
76                 fmt.Fprintln(stdout, output+"/dumpgob.txt")
77                 return 0
78         }
79
80         input, err := open(*inputFilename)
81         if err != nil {
82                 return 1
83         }
84         defer input.Close()
85         output, err := os.OpenFile(*outputFilename, os.O_CREATE|os.O_WRONLY, 0644)
86         if err != nil {
87                 return 1
88         }
89         defer output.Close()
90         bufw := bufio.NewWriterSize(output, 8*1024*1024)
91
92         var n, nCG, nCS, nTV int
93         err = DecodeLibrary(input, strings.HasSuffix(*inputFilename, ".gz"), func(ent *LibraryEntry) error {
94                 if n%1000000 == 0 {
95                         fmt.Fprintf(stderr, "ent %d\n", n)
96                 }
97                 n++
98                 if len(ent.TagSet) > 0 {
99                         fmt.Fprintf(bufw, "ent %d: TagSet, len %d, taglen %d\n", n, len(ent.TagSet), len(ent.TagSet[0]))
100                 }
101                 for _, cg := range ent.CompactGenomes {
102                         nCG++
103                         fmt.Fprintf(bufw, "ent %d: CompactGenome, name %q, len(Variants) %d\n", n, cg.Name, len(cg.Variants))
104                 }
105                 for _, cs := range ent.CompactSequences {
106                         nCS++
107                         fmt.Fprintf(bufw, "ent %d: CompactSequence, name %q, len(TileSequences) %d\n", n, cs.Name, len(cs.TileSequences))
108                 }
109                 for _, tv := range ent.TileVariants {
110                         nTV++
111                         fmt.Fprintf(bufw, "ent %d: TileVariant, tag %d, variant %d, hash %x, len(seq) %d\n", n, tv.Tag, tv.Variant, tv.Blake2b, len(tv.Sequence))
112                 }
113                 return nil
114         })
115         if err != nil {
116                 return 1
117         }
118         fmt.Fprintf(bufw, "total: ents %d, CompactGenomes %d, CompactSequences %d, TileVariants %d\n", n, nCG, nCS, nTV)
119         err = bufw.Flush()
120         if err != nil {
121                 return 1
122         }
123         err = output.Close()
124         if err != nil {
125                 return 1
126         }
127         return 0
128 }