Merge pull request pywren#88 from pywren/sharding-config

Move runtime sharding into config variable

Author: ooq

pywren/executor.py

@@ -8,6 +8,7 @@
     import pickle
 from multiprocessing.pool import ThreadPool
 import time
+import random
 import logging
 import botocore
 import glob2
@@ -31,7 +32,7 @@ class Executor(object):
     """
 
     def __init__(self, aws_region, s3_bucket, s3_prefix,
-                 invoker, config, job_max_runtime, shard_runtime=False):
+                 invoker, config, job_max_runtime):
         self.aws_region = aws_region
         self.s3_bucket = s3_bucket
         self.s3_prefix = s3_prefix
@@ -41,11 +42,11 @@ def __init__(self, aws_region, s3_bucket, s3_prefix,
         self.invoker = invoker
         self.s3client = self.session.create_client('s3', region_name = aws_region)
         self.job_max_runtime = job_max_runtime
-        self.shard_runtime = shard_runtime
 
         runtime_bucket = config['runtime']['s3_bucket']
         runtime_key =  config['runtime']['s3_key']
-        if not runtime.runtime_key_valid(runtime_bucket, runtime_key):
+        self.runtime_meta_info = runtime.get_runtime_info(runtime_bucket, runtime_key)
+        if not runtime.runtime_key_valid(self.runtime_meta_info):
             raise Exception("The indicated runtime: s3://{}/{} is not approprite for this python version".format(runtime_bucket, runtime_key))
 
     def create_mod_data(self, mod_paths):
@@ -85,6 +86,18 @@ def invoke_with_keys(self, s3_func_key, s3_data_key, s3_output_key,
                          host_job_meta, job_max_runtime,
                          overwrite_invoke_args = None):
 
+        # Pick a runtime url if we have shards. If not the handler will construct it
+        # using s3_bucket and s3_key
+        runtime_url = ""
+        if ('urls' in self.runtime_meta_info and
+                isinstance(self.runtime_meta_info['urls'], list) and
+                    len(self.runtime_meta_info['urls']) > 1):
+            num_shards = len(self.runtime_meta_info['urls'])
+            logger.debug("Runtime is sharded, choosing from {} copies.".format(num_shards))
+            random.seed()
+            runtime_url = random.choice(self.runtime_meta_info['urls'])
+
+
         arg_dict = {'func_key' : s3_func_key,
                     'data_key' : s3_data_key,
                     'output_key' : s3_output_key,
@@ -97,7 +110,7 @@ def invoke_with_keys(self, s3_func_key, s3_data_key, s3_output_key,
                     'runtime_s3_bucket' : self.config['runtime']['s3_bucket'],
                     'runtime_s3_key' : self.config['runtime']['s3_key'],
                     'pywren_version' : version.__version__,
-                    'shard_runtime_key' : self.shard_runtime}
+                    'runtime_url' : runtime_url }
 
         if extra_env is not None:
             logger.debug("Extra environment vars {}".format(extra_env))

pywren/runtime.py

@@ -18,8 +18,7 @@ def get_runtime_info(bucket, key):
 def version_str(version_info):
     return "{}.{}".format(version_info[0], version_info[1])
 
-def runtime_key_valid(bucket, key):
-    runtime_meta = get_runtime_info(bucket, key)
+def runtime_key_valid(runtime_meta):
     return runtime_valid(runtime_meta)
 
 def runtime_valid(runtime_meta):

pywren/wren.py

@@ -6,6 +6,9 @@
     import pickle
 from tblib import pickling_support
 import logging
+import botocore
+import glob2
+import random
 import os
 pickling_support.install()
 
@@ -32,7 +35,7 @@ def default_executor(**kwargs):
     return lambda_executor(**kwargs)
 
 
-def lambda_executor(config= None, job_max_runtime=280, shard_runtime=False):
+def lambda_executor(config=None, job_max_runtime=280):
 
     if config is None:
         config = wrenconfig.default()
@@ -43,22 +46,21 @@ def lambda_executor(config= None, job_max_runtime=280, shard_runtime=False):
     S3_PREFIX = config['s3']['pywren_prefix']
 
     invoker = invokers.LambdaInvoker(AWS_REGION, FUNCTION_NAME)
-    return Executor(AWS_REGION, S3_BUCKET, S3_PREFIX, invoker, config,
-                    job_max_runtime, shard_runtime=shard_runtime)
+    return Executor(AWS_REGION, S3_BUCKET, S3_PREFIX, invoker, config, 
+                    job_max_runtime)
 
 
-def dummy_executor(shard_runtime=False):
+def dummy_executor():
     config = wrenconfig.default()
     AWS_REGION = config['account']['aws_region']
     S3_BUCKET = config['s3']['bucket']
     S3_PREFIX = config['s3']['pywren_prefix']
     invoker = invokers.DummyInvoker()
     return Executor(AWS_REGION, S3_BUCKET, S3_PREFIX, invoker, config,
-                    100, shard_runtime=shard_runtime)
+                    100)
 
 
-def remote_executor(config= None, job_max_runtime=3600,
-                    shard_runtime=False):
+def remote_executor(config= None, job_max_runtime=3600):
     if config is None:
         config = wrenconfig.default()
 
@@ -68,7 +70,9 @@ def remote_executor(config= None, job_max_runtime=3600,
     S3_PREFIX = config['s3']['pywren_prefix']
     invoker = invokers.SQSInvoker(AWS_REGION, SQS_QUEUE)
     return Executor(AWS_REGION, S3_BUCKET, S3_PREFIX, invoker, config,
-                    job_max_runtime, shard_runtime=shard_runtime)
+                    job_max_runtime)
+
+standalone_executor = remote_executor
 
 
 def get_all_results(fs):

pywren/wrenconfig.py

@@ -21,7 +21,6 @@
 
 MAX_AGG_DATA_SIZE = 4e6
 
-MAX_S3_RUNTIME_SHARDS=50
 default_runtime = {'2.7' : "pywren.runtime/pywren_runtime-2.7-default.tar.gz", 
                    '3.5' : "pywren.runtime/pywren_runtime-3.5-default.tar.gz", 
                    '3.6' : "pywren.runtime/pywren_runtime-3.6-default.tar.gz"}

pywren/wrenhandler.py

@@ -13,22 +13,17 @@
 import traceback
 from threading import Thread
 import signal
-import random
 
 if (sys.version_info > (3, 0)):
     from . import wrenutil
     from . import s3util
     from . import version
-    from . import wrenconfig
-    from . import wrenlogging
     from queue import Queue, Empty
 
 else:
     import wrenutil
     import s3util
     import version
-    import wrenconfig
-    import wrenlogging
     from Queue import Queue, Empty
 
 
@@ -145,12 +140,11 @@ def generic_handler(event, context_dict):
 
         runtime_s3_bucket = event['runtime_s3_bucket']
         runtime_s3_key = event['runtime_s3_key']
-        if event.get('shard_runtime_key', False):
-            random.seed()
-            shard = random.randrange(wrenconfig.MAX_S3_RUNTIME_SHARDS)
-            key_shard = wrenutil.get_s3_shard(runtime_s3_key, shard)
-            runtime_s3_key_used = wrenutil.hash_s3_key(key_shard)
+        if event.get('runtime_url'):
+            # NOTE(shivaram): Right now we only support S3 urls.
+            runtime_s3_bucket_used, runtime_s3_key_used = wrenutil.split_s3_url(event['runtime_url'])
         else:
+            runtime_s3_bucket_used = runtime_s3_bucket
             runtime_s3_key_used = runtime_s3_key
 
         job_max_runtime = event.get("job_max_runtime", 290) # default for lambda
@@ -224,8 +218,9 @@ def generic_handler(event, context_dict):
         logger.debug(subprocess.check_output("find {}".format(os.getcwd()), shell=True))
 
         response_status['runtime_s3_key_used'] = runtime_s3_key_used
+        response_status['runtime_s3_bucket_used'] = runtime_s3_bucket_used
 
-        runtime_cached = download_runtime_if_necessary(s3, runtime_s3_bucket, 
+        runtime_cached = download_runtime_if_necessary(s3, runtime_s3_bucket_used,
                                                        runtime_s3_key_used)
         logger.info("Runtime ready, cached={}".format(runtime_cached))
         response_status['runtime_cached'] = runtime_cached

pywren/wrenutil.py

@@ -117,17 +117,3 @@ def split_s3_url(s3_url):
     bucket_name = splits[0]
     key = "/".join(splits[1:])
     return bucket_name, key
-
-def hash_s3_key(s):
-    """
-    MD5-hash the contents of an S3 key to enable good partitioning. 
-    used for sharding the runtimes
-    """
-    DIGEST_LEN = 6
-    m = hashlib.md5()
-    m.update(s.encode('ascii'))
-    digest = m.hexdigest()
-    return "{}-{}".format(digest[:DIGEST_LEN], s)
-
-def get_s3_shard(key, shard_num):
-    return "{}.{:04d}".format(key, shard_num)

tests/test_simple.py

@@ -4,6 +4,7 @@
 import uuid
 import numpy as np
 import time
+import os
 import pywren
 import pywren.runtime
 import subprocess
@@ -283,19 +284,29 @@ class RuntimeSharding(unittest.TestCase):
     download the real key
     """
     def test_no_shard(self):
-        wrenexec = pywren.default_executor(shard_runtime=False)
+        config = pywren.wrenconfig.default()
+        old_key = config['runtime']['s3_key']
+        prefix, tar_gz = os.path.split(old_key)
+        # Use the staging key to test as it doesn't have shards
+        config['runtime']['s3_key'] = os.path.join("pywren.runtime.staging", tar_gz)
+        wrenexec = pywren.default_executor(config=config)
 
         def test_func(x):
             return x + 1
 
         future = wrenexec.call_async(test_func, 7)
         result = future.result()
         base_runtime_key = wrenexec.config['runtime']['s3_key']
-        self.assertEqual(future.run_status['runtime_s3_key_used'], 
+        self.assertEqual(future.run_status['runtime_s3_key_used'],
                          base_runtime_key)
 
     def test_shard(self):
-        wrenexec = pywren.default_executor(shard_runtime=True)
+        config = pywren.wrenconfig.default()
+        old_key = config['runtime']['s3_key']
+        prefix, tar_gz = os.path.split(old_key)
+        # Use a runtime that has shards
+        config['runtime']['s3_key'] = os.path.join("pywren.runtime", tar_gz)
+        wrenexec = pywren.default_executor(config=config)
 
         def test_func(x):
             return x + 1
@@ -304,5 +315,6 @@ def test_func(x):
 
         future = wrenexec.call_async(test_func, 7)
         result = future.result()
+        # NOTE: There is some probability we will hit the base key ? 
         self.assertNotEqual(future.run_status['runtime_s3_key_used'], 
                          base_runtime_key)

tests/test_util.py

@@ -16,11 +16,3 @@ def test_s3_split(self):
         with pytest.raises(ValueError) as excinfo:
             bad_s3_url = "notS3://foo/bar"
             bucket, key = pywren.wrenutil.split_s3_url(bad_s3_url)
-            
-    def test_hash(self):
-        
-        key = 'testkey'
-        hashed_key = pywren.wrenutil.hash_s3_key(key)
-        print(hashed_key)
-        self.assertEqual(hashed_key[-len(key):], key)
-        self.assertNotEqual(hashed_key, key)