Skip to content

Commit 05821ea

Browse files
authored
Merge pull request #74 from worldcoin/piohei/add_start_from_s3
Add start-from-s3 command.
2 parents 1daf57f + 2f53532 commit 05821ea

File tree

5 files changed

+179
-3
lines changed

5 files changed

+179
-3
lines changed

go.mod

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,10 @@ go 1.23
55
toolchain go1.23.3
66

77
require (
8+
github.com/aws/aws-sdk-go-v2 v1.33.0
9+
github.com/aws/aws-sdk-go-v2/config v1.29.1
10+
github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.17.53
11+
github.com/aws/aws-sdk-go-v2/service/s3 v1.74.0
812
github.com/consensys/gnark v0.8.0
913
github.com/iden3/go-iden3-crypto v0.0.13
1014
github.com/prometheus/client_golang v1.14.0
@@ -22,6 +26,21 @@ require (
2226
github.com/DataDog/go-tuf v1.1.0-0.5.2 // indirect
2327
github.com/DataDog/sketches-go v1.4.5 // indirect
2428
github.com/Microsoft/go-winio v0.6.1 // indirect
29+
github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.7 // indirect
30+
github.com/aws/aws-sdk-go-v2/credentials v1.17.54 // indirect
31+
github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.24 // indirect
32+
github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.28 // indirect
33+
github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.28 // indirect
34+
github.com/aws/aws-sdk-go-v2/internal/ini v1.8.1 // indirect
35+
github.com/aws/aws-sdk-go-v2/internal/v4a v1.3.28 // indirect
36+
github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.12.1 // indirect
37+
github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.5.2 // indirect
38+
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.12.9 // indirect
39+
github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.18.9 // indirect
40+
github.com/aws/aws-sdk-go-v2/service/sso v1.24.11 // indirect
41+
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.28.10 // indirect
42+
github.com/aws/aws-sdk-go-v2/service/sts v1.33.9 // indirect
43+
github.com/aws/smithy-go v1.22.1 // indirect
2544
github.com/beorn7/perks v1.0.1 // indirect
2645
github.com/blang/semver/v4 v4.0.0 // indirect
2746
github.com/cespare/xxhash/v2 v2.2.0 // indirect

go.sum

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,44 @@ github.com/alecthomas/units v0.0.0-20151022065526-2efee857e7cf/go.mod h1:ybxpYRF
5858
github.com/alecthomas/units v0.0.0-20190717042225-c3de453c63f4/go.mod h1:ybxpYRFXyAe+OPACYpWeL0wqObRcbAqCMya13uyzqw0=
5959
github.com/alecthomas/units v0.0.0-20190924025748-f65c72e2690d/go.mod h1:rBZYJk541a8SKzHPHnH3zbiI+7dagKZ0cgpgrD7Fyho=
6060
github.com/armon/go-radix v0.0.0-20180808171621-7fddfc383310/go.mod h1:ufUuZ+zHj4x4TnLV4JWEpy2hxWSpsRywHrMgIH9cCH8=
61+
github.com/aws/aws-sdk-go-v2 v1.33.0 h1:Evgm4DI9imD81V0WwD+TN4DCwjUMdc94TrduMLbgZJs=
62+
github.com/aws/aws-sdk-go-v2 v1.33.0/go.mod h1:P5WJBrYqqbWVaOxgH0X/FYYD47/nooaPOZPlQdmiN2U=
63+
github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.7 h1:lL7IfaFzngfx0ZwUGOZdsFFnQ5uLvR0hWqqhyE7Q9M8=
64+
github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.7/go.mod h1:QraP0UcVlQJsmHfioCrveWOC1nbiWUl3ej08h4mXWoc=
65+
github.com/aws/aws-sdk-go-v2/config v1.29.1 h1:JZhGawAyZ/EuJeBtbQYnaoftczcb2drR2Iq36Wgz4sQ=
66+
github.com/aws/aws-sdk-go-v2/config v1.29.1/go.mod h1:7bR2YD5euaxBhzt2y/oDkt3uNRb6tjFp98GlTFueRwk=
67+
github.com/aws/aws-sdk-go-v2/credentials v1.17.54 h1:4UmqeOqJPvdvASZWrKlhzpRahAulBfyTJQUaYy4+hEI=
68+
github.com/aws/aws-sdk-go-v2/credentials v1.17.54/go.mod h1:RTdfo0P0hbbTxIhmQrOsC/PquBZGabEPnCaxxKRPSnI=
69+
github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.24 h1:5grmdTdMsovn9kPZPI23Hhvp0ZyNm5cRO+IZFIYiAfw=
70+
github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.24/go.mod h1:zqi7TVKTswH3Ozq28PkmBmgzG1tona7mo9G2IJg4Cis=
71+
github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.17.53 h1:3jYpOndmkKtmlPOhMNIV7Q92GD61x/KNjmxUcB95btw=
72+
github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.17.53/go.mod h1:+s7tPUl4uy7FMpT5qnjkY5YJNuKU2HZL6trkYxQNtb4=
73+
github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.28 h1:igORFSiH3bfq4lxKFkTSYDhJEUCYo6C8VKiWJjYwQuQ=
74+
github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.28/go.mod h1:3So8EA/aAYm36L7XIvCVwLa0s5N0P7o2b1oqnx/2R4g=
75+
github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.28 h1:1mOW9zAUMhTSrMDssEHS/ajx8JcAj/IcftzcmNlmVLI=
76+
github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.28/go.mod h1:kGlXVIWDfvt2Ox5zEaNglmq0hXPHgQFNMix33Tw22jA=
77+
github.com/aws/aws-sdk-go-v2/internal/ini v1.8.1 h1:VaRN3TlFdd6KxX1x3ILT5ynH6HvKgqdiXoTxAF4HQcQ=
78+
github.com/aws/aws-sdk-go-v2/internal/ini v1.8.1/go.mod h1:FbtygfRFze9usAadmnGJNc8KsP346kEe+y2/oyhGAGc=
79+
github.com/aws/aws-sdk-go-v2/internal/v4a v1.3.28 h1:7kpeALOUeThs2kEjlAxlADAVfxKmkYAedlpZ3kdoSJ4=
80+
github.com/aws/aws-sdk-go-v2/internal/v4a v1.3.28/go.mod h1:pyaOYEdp1MJWgtXLy6q80r3DhsVdOIOZNB9hdTcJIvI=
81+
github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.12.1 h1:iXtILhvDxB6kPvEXgsDhGaZCSC6LQET5ZHSdJozeI0Y=
82+
github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.12.1/go.mod h1:9nu0fVANtYiAePIBh2/pFUSwtJ402hLnp854CNoDOeE=
83+
github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.5.2 h1:e6um6+DWYQP1XCa+E9YVtG/9v1qk5lyAOelMOVwSyO8=
84+
github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.5.2/go.mod h1:dIW8puxSbYLSPv/ju0d9A3CpwXdtqvJtYKDMVmPLOWE=
85+
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.12.9 h1:TQmKDyETFGiXVhZfQ/I0cCFziqqX58pi4tKJGYGFSz0=
86+
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.12.9/go.mod h1:HVLPK2iHQBUx7HfZeOQSEu3v2ubZaAY2YPbAm5/WUyY=
87+
github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.18.9 h1:2aInXbh02XsbO0KobPGMNXyv2QP73VDKsWPNJARj/+4=
88+
github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.18.9/go.mod h1:dgXS1i+HgWnYkPXqNoPIPKeUsUUYHaUbThC90aDnNiE=
89+
github.com/aws/aws-sdk-go-v2/service/s3 v1.74.0 h1:ncCHiFU9Eq4qnKCNlzMZXfFmvb9R8OVNfU8SFOskxdI=
90+
github.com/aws/aws-sdk-go-v2/service/s3 v1.74.0/go.mod h1:jGJ/v7FIi7Ys9t54tmEFnrxuaWeJLpwNgKp2DXAVhOU=
91+
github.com/aws/aws-sdk-go-v2/service/sso v1.24.11 h1:kuIyu4fTT38Kj7YCC7ouNbVZSSpqkZ+LzIfhCr6Dg+I=
92+
github.com/aws/aws-sdk-go-v2/service/sso v1.24.11/go.mod h1:Ro744S4fKiCCuZECXgOi760TiYylUM8ZBf6OGiZzJtY=
93+
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.28.10 h1:l+dgv/64iVlQ3WsBbnn+JSbkj01jIi+SM0wYsj3y/hY=
94+
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.28.10/go.mod h1:Fzsj6lZEb8AkTE5S68OhcbBqeWPsR8RnGuKPr8Todl8=
95+
github.com/aws/aws-sdk-go-v2/service/sts v1.33.9 h1:BRVDbewN6VZcwr+FBOszDKvYeXY1kJ+GGMCcpghlw0U=
96+
github.com/aws/aws-sdk-go-v2/service/sts v1.33.9/go.mod h1:f6vjfZER1M17Fokn0IzssOTMT2N8ZSq+7jnNF0tArvw=
97+
github.com/aws/smithy-go v1.22.1 h1:/HPHZQ0g7f4eUeK6HKglFz8uwVfZKgoI25rb/J+dnro=
98+
github.com/aws/smithy-go v1.22.1/go.mod h1:irrKGvNn1InZwb2d7fkIRNucdfwR8R+Ts3wxYa/cJHg=
6199
github.com/beorn7/perks v0.0.0-20180321164747-3a771d992973/go.mod h1:Dwedo/Wpr24TaqPxmxbtue+5NUziq4I4S80YR8gNf3Q=
62100
github.com/beorn7/perks v1.0.0/go.mod h1:KWe93zE9D1o94FZ5RNwFwVgaQK1VOXiVxmqh+CedLV8=
63101
github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=

main.go

Lines changed: 70 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import (
1010
"time"
1111

1212
"worldcoin/gnark-mbu/logging"
13+
"worldcoin/gnark-mbu/poseidon_tree"
1314
"worldcoin/gnark-mbu/prover"
1415
"worldcoin/gnark-mbu/server"
1516

@@ -226,7 +227,7 @@ func main() {
226227

227228
if mode == server.InsertionMode {
228229
params := prover.InsertionParameters{}
229-
tree := NewTree(treeDepth)
230+
tree := poseidon_tree.NewTree(treeDepth)
230231

231232
params.StartIndex = 0
232233
params.PreRoot = tree.Root()
@@ -241,7 +242,7 @@ func main() {
241242
r, err = json.Marshal(&params)
242243
} else if mode == server.DeletionMode {
243244
params := prover.DeletionParameters{}
244-
tree := NewTree(treeDepth)
245+
tree := poseidon_tree.NewTree(treeDepth)
245246

246247
params.DeletionIndices = make([]uint32, batchSize)
247248
params.IdComms = make([]big.Int, batchSize)
@@ -320,6 +321,73 @@ func main() {
320321
return nil
321322
},
322323
},
324+
{
325+
Name: "start-from-s3",
326+
Flags: []cli.Flag{
327+
&cli.StringFlag{Name: "mode", Usage: "insertion/deletion", EnvVars: []string{"MTB_MODE"}, DefaultText: "insertion"},
328+
&cli.BoolFlag{Name: "json-logging", Usage: "enable JSON logging", Required: false},
329+
&cli.StringFlag{Name: "prover-address", Usage: "address for the prover server", Value: "localhost:3001", Required: false},
330+
&cli.StringFlag{Name: "metrics-address", Usage: "address for the metrics server", Value: "localhost:9998", Required: false},
331+
&cli.StringFlag{Name: "s3-region", Usage: "s3 region of bucket", EnvVars: []string{"S3_REGION"}, DefaultText: "us-east1"},
332+
&cli.StringFlag{Name: "s3-bucket", Usage: "s3 bucket name", EnvVars: []string{"S3_BUCKET"}, Required: true},
333+
&cli.StringFlag{Name: "s3-object-key", Usage: "s3 object key (path)", EnvVars: []string{"S3_OBJECT_KEY"}, Required: true},
334+
&cli.IntFlag{Name: "s3-concurrency", Usage: "number of concurrent connections to download from s3", EnvVars: []string{"S3_CONCURRENCY"}, DefaultText: "8"},
335+
&cli.Int64Flag{Name: "s3-part-mibs", Usage: "size of part to download from s3", EnvVars: []string{"S3_PART_MIBS"}, DefaultText: "64"},
336+
},
337+
Action: func(context *cli.Context) error {
338+
if context.Bool("json-logging") {
339+
logging.SetJSONOutput()
340+
}
341+
region := context.String("s3-region")
342+
bucket := context.String("s3-bucket")
343+
objectKey := context.String("s3-object-key")
344+
concurrency := context.Int("s3-concurrency")
345+
partMibs := context.Int64("s3-part-mibs")
346+
mode := context.String("mode")
347+
348+
if mode != server.DeletionMode && mode != server.InsertionMode {
349+
return fmt.Errorf("invalid mode: %s", mode)
350+
}
351+
352+
logging.Logger().
353+
Info().
354+
Str("region", region).
355+
Str("bucket", bucket).
356+
Str("objectKey", objectKey).
357+
Str("objectKey", objectKey).
358+
Int("concurrency", concurrency).
359+
Int64("partMibs", partMibs).
360+
Msg("Loading proving system from S3")
361+
362+
start := time.Now()
363+
ps, err := prover.ReadSystemFromS3(region, bucket, objectKey, concurrency, partMibs)
364+
if err != nil {
365+
return err
366+
}
367+
duration := time.Since(start)
368+
logging.Logger().
369+
Info().
370+
Uint32("treeDepth", ps.TreeDepth).
371+
Uint32("batchSize", ps.BatchSize).
372+
Dur("duration", duration).
373+
Msg("Proving system loaded")
374+
375+
config := server.Config{
376+
ProverAddress: context.String("prover-address"),
377+
MetricsAddress: context.String("metrics-address"),
378+
Mode: mode,
379+
}
380+
instance := server.Run(&config, ps)
381+
sigint := make(chan os.Signal, 1)
382+
signal.Notify(sigint, os.Interrupt)
383+
<-sigint
384+
logging.Logger().Info().Msg("Received sigint, shutting down")
385+
instance.RequestStop()
386+
logging.Logger().Info().Msg("Waiting for server to close")
387+
instance.AwaitStop()
388+
return nil
389+
},
390+
},
323391
{
324392
Name: "prove",
325393
Flags: []cli.Flag{

test_tree.go renamed to poseidon_tree/poseidon_tree.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
package main
1+
package poseidon_tree
22

33
import (
44
"github.com/iden3/go-iden3-crypto/poseidon"

prover/marshal.go

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,15 +3,20 @@ package prover
33
import (
44
"bufio"
55
"bytes"
6+
"context"
67
"encoding/binary"
78
"encoding/json"
89
"fmt"
10+
"github.com/aws/aws-sdk-go-v2/config"
911
"io"
1012
"math/big"
1113
"os"
1214

1315
"worldcoin/gnark-mbu/logging"
1416

17+
"github.com/aws/aws-sdk-go-v2/aws"
18+
"github.com/aws/aws-sdk-go-v2/feature/s3/manager"
19+
"github.com/aws/aws-sdk-go-v2/service/s3"
1520
"github.com/consensys/gnark-crypto/ecc"
1621
"github.com/consensys/gnark/backend/groth16"
1722
)
@@ -392,3 +397,49 @@ func ReadSystemFromFile(path string) (ps *ProvingSystem, err error) {
392397
}
393398
return
394399
}
400+
401+
func ReadSystemFromS3(region, bucket, objectKey string, concurrency int, partMiBs int64) (ps *ProvingSystem, err error) {
402+
ps = new(ProvingSystem)
403+
404+
ctx := context.TODO()
405+
406+
awsConfig, err := config.LoadDefaultConfig(ctx, config.WithRegion(region))
407+
if err != nil {
408+
return
409+
}
410+
411+
client := s3.NewFromConfig(awsConfig)
412+
if err != nil {
413+
return
414+
}
415+
416+
hObj, err := client.HeadObject(ctx, &s3.HeadObjectInput{
417+
Bucket: aws.String(bucket),
418+
Key: aws.String(objectKey),
419+
})
420+
if err != nil {
421+
return
422+
}
423+
424+
downloader := manager.NewDownloader(client, func(d *manager.Downloader) {
425+
d.PartSize = partMiBs * 1024 * 1024
426+
d.Concurrency = concurrency
427+
})
428+
429+
buff := make([]byte, *hObj.ContentLength)
430+
w := manager.NewWriteAtBuffer(buff)
431+
_, err = downloader.Download(context.TODO(), w, &s3.GetObjectInput{
432+
Bucket: aws.String(bucket),
433+
Key: aws.String(objectKey),
434+
})
435+
if err != nil {
436+
return
437+
}
438+
439+
bufferedReader := bytes.NewReader(w.Bytes())
440+
_, err = ps.UnsafeReadFrom(bufferedReader)
441+
if err != nil {
442+
return
443+
}
444+
return
445+
}

0 commit comments

Comments
 (0)