[tor-commits] [sbws/master] Add args to generate to scale as torflow too

[juga at torproject.org]

commit 028dc8bff6dd48a6b5a0682ba4904b70532842f0
Author: juga0 <juga at riseup.net>
Date:   Fri Aug 31 11:01:25 2018 +0000

    Add args to generate to scale as torflow too
---
 sbws/core/generate.py | 17 ++++++++++++++---
 1 file changed, 14 insertions(+), 3 deletions(-)

diff --git a/sbws/core/generate.py b/sbws/core/generate.py
index 7d68be7..f996c93 100644
--- a/sbws/core/generate.py
+++ b/sbws/core/generate.py
@@ -1,4 +1,5 @@
-from sbws.globals import (fail_hard, SBWS_SCALE_CONSTANT)
+from sbws.globals import (fail_hard, SBWS_SCALE_CONSTANT, TORFLOW_SCALING,
+                          SBWS_SCALING)
 from sbws.lib.v3bwfile import V3BWFile
 from sbws.lib.resultdump import load_recent_results_in_datadir
 from argparse import ArgumentDefaultsHelpFormatter
@@ -31,11 +32,14 @@ def gen_parser(sub):
     p.add_argument('--scale-constant', default=SBWS_SCALE_CONSTANT, type=int,
                    help='When scaling bw weights, scale them using this const '
                    'multiplied by the number of measured relays')
-    p.add_argument('--scale', action='store_true',
+    p.add_argument('--scale_sbws', action='store_true',
                    help='If specified, do not use bandwidth values as they '
                    'are, but scale them such that we have a budget of '
                    'scale_constant * num_measured_relays = bandwidth to give '
                    'out, and we do so proportionally')
+    p.add_argument('-t', '--scale_torflow', action='store_true',
+                   help='If specified, do not use bandwidth values as they '
+                   'are, but scale them in the way Torflow does.')
 
 
 def main(args, conf):
@@ -46,6 +50,12 @@ def main(args, conf):
         fail_hard('%s does not exist', datadir)
     if args.scale_constant < 1:
         fail_hard('--scale-constant must be positive')
+    if args.scale_sbws:
+        scaling_method = SBWS_SCALING
+    elif args.scale_torflow:
+        scaling_method = TORFLOW_SCALING
+    else:
+        scaling_method = None
 
     fresh_days = conf.getint('general', 'data_period')
     reset_bw_ipv4_changes = conf.getboolean('general', 'reset_bw_ipv4_changes')
@@ -59,7 +69,8 @@ def main(args, conf):
                     'ran sbws scanner recently?)')
         return
     state_fpath = conf.getpath('paths', 'state_fname')
-    bw_file = V3BWFile.from_results(results, state_fpath, args.scale_constant)
+    bw_file = V3BWFile.from_results(results, state_fpath, args.scale_constant,
+                                    scaling_method)
     output = args.output or \
         conf.getpath('paths', 'v3bw_fname').format(now_fname())
     bw_file.write(output)