Repository: sosw/sosw

Issue:
Scheduler: Get different job_schemas per invocation
Should support providing custom job_schemas for different calls.
Options:
 - Have multiple preconfigured ones and automatically find first suitable (probably bad idea)
 - Provide optional custom job_schema in the payload (already supported, but tricky to call).
 - Provide `class` of job in the invocation and take the specific job_schema from DynamoDB config ON CALL (smart but more work).
 - Other ideas...

Discussion:

Patch:
diff --git a/.gitignore b/.gitignore
index 1fa2656..9f0819d 100644
--- a/.gitignore
+++ b/.gitignore
@@ -65,6 +65,7 @@ instance/
 
 # Sphinx documentation
 docs/_build/
+sosw-rtd/*
 
 # PyBuilder
 target/
diff --git a/docs/essentials/scheduler.rst b/docs/essentials/scheduler.rst
index b3091e2..16e402c 100644
--- a/docs/essentials/scheduler.rst
+++ b/docs/essentials/scheduler.rst
@@ -1,5 +1,6 @@
 .. _Scheduler:
 
+---------
 Scheduler
 ---------
 
@@ -14,55 +15,69 @@ configurable rules for chunking specific for different workers.
    Scheduler Workflow
 
 
-.. code-block:: python
+..  code-block:: python
+
+    TASKS_TABLE_CONFIG = {
+        'row_mapper':       {
+            'task_id':             'S',
+            'labourer_id':         'S',
+            'greenfield':          'N',
+            'attempts':            'N',
+            'closed_at':           'N',
+            'completed_at':        'N',
+            'desired_launch_time': 'N',
+            'arn':                 'S',
+            'payload':             'S'
+        },
+        'required_fields':  ['task_id', 'labourer_id', 'created_at', 'greenfield'],
+        'table_name':       'sosw_tasks',
+        'index_greenfield': 'sosw_tasks_greenfield',
+        'field_names':      {
+            'task_id':     'task_id',
+            'labourer_id': 'labourer_id',
+            'greenfield':  'greenfield',
+        }
+    }
+
+    TASK_CLIENT_CONFIG = {
+        'dynamo_db_config':                  TASKS_TABLE_CONFIG,
+        'sosw_closed_tasks_table':           'sosw_closed_tasks',
+        'sosw_retry_tasks_table':            'sosw_retry_tasks',
+        'sosw_retry_tasks_greenfield_index': 'labourer_id_greenfield',
+        'ecology_config':             {},
+        'labourers':                         {
+            'some_function': {
+                'arn':                          'arn:aws:lambda:us-west-2:0000000000:function:some_function',
+                'max_simultaneous_invocations': 10,
+            },
+        },
+    }
 
-   TASKS_TABLE_CONFIG = {
-       'row_mapper':       {
-           'task_id':             'S',
-           'labourer_id':         'S',
-           'greenfield':          'N',
-           'attempts':            'N',
-           'closed_at':           'N',
-           'completed_at':        'N',
-           'desired_launch_time': 'N',
-           'arn':                 'S',
-           'payload':             'S'
-       },
-       'required_fields':  ['task_id', 'labourer_id', 'created_at', 'greenfield'],
-       'table_name':       'sosw_tasks',
-       'index_greenfield': 'sosw_tasks_greenfield',
-       'field_names':      {
-           'task_id':     'task_id',
-           'labourer_id': 'labourer_id',
-           'greenfield':  'greenfield',
-       }
-   }
+    SCHEDULER_CONFIG = {
 
-   TASK_CLIENT_CONFIG = {
-       'dynamo_db_config':                  TASKS_TABLE_CONFIG,
-       'sosw_closed_tasks_table':           'sosw_closed_tasks',
-       'sosw_retry_tasks_table':            'sosw_retry_tasks',
-       'sosw_retry_tasks_greenfield_index': 'labourer_id_greenfield',
-       'ecology_config':             {},
-       'labourers':                         {
-           'some_function': {
-               'arn':                          'arn:aws:lambda:us-west-2:0000000000:function:some_function',
-               'max_simultaneous_invocations': 10,
-           },
-       },
-   }
+        'queue_bucket':       'some-bucket',
 
-   SCHEDULER_CONFIG = {
+        'task_config': TASK_CLIENT_CONFIG,
+        'job_schema': {},
+        'job_schema_variants': {
+            'default': {
+                'chunkable_attrs': [
+                    # ('section', {}),
+                    # ('store', {}),
+                    # ('product', {}),
+                ]
+            }
+        }
+    }
 
-       'queue_bucket':       'some-bucket',
 
-       'task_config': TASK_CLIENT_CONFIG,
-       'job_schema':         {
-           'chunkable_attrs': [
-           ]
-       }
-   }
+Job Schema
+----------
 
+After the Scheduler has parsed and validated the job from the event, it provides `custom_config` from this job.
+In case the `job_schema_name` is passed, a certain `job_schema` type will be applied from `job_schema_variants`,
+otherwise it takes the default one.
+Each time the Scheduler is called it overwrites the `custom_config` and use new specified `job_schema` type.
 
 .. automodule:: sosw.scheduler
    :members:
diff --git a/sosw/scheduler.py b/sosw/scheduler.py
index 647673c..dd3dbee 100644
--- a/sosw/scheduler.py
+++ b/sosw/scheduler.py
@@ -5,7 +5,7 @@
     sosw - Serverless Orchestrator of Serverless Workers
 
     The MIT License (MIT)
-    Copyright (C) 2020  sosw core contributors <info@sosw.app>
+    Copyright (C) 2021  sosw core contributors <info@sosw.app>
 
     Permission is hereby granted, free of charge, to any person obtaining a copy
     of this software and associated documentation files (the "Software"), to deal
@@ -97,12 +97,15 @@ class Scheduler(Essential):
         'queue_bucket':    'autotest-bucket',
         'shutdown_period': 60,
         'rows_to_process': 50,
-        'job_schema':      {
-            'chunkable_attrs': [
-                # ('section', {}),
-                # ('store', {}),
-                # ('product', {}),
-            ],
+        'job_schema':      {},
+        'job_schema_variants': {
+            'default': {
+                'chunkable_attrs': [
+                    # ('section', {}),
+                    # ('store', {}),
+                    # ('product', {}),
+                ]
+            }
         },
         'task_operational_overhead_for_ddb': 0.03,
     }
@@ -135,6 +138,7 @@ class Scheduler(Essential):
         """
 
         job = self.extract_job_from_payload(event)
+        self.apply_job_schema(name=job.get('job_schema_name'))
 
         # If called as sibling
         if 'file_name' in job:
@@ -149,6 +153,15 @@ class Scheduler(Essential):
         super().__call__(event)
 
 
+    def apply_job_schema(self, name: str = None):
+        """
+        Apply a job_schema from job_schema_variants by the name or apply the default one.
+
+        """
+
+        self.config['job_schema'] = self.config['job_schema_variants'][name or 'default']
+
+
     def parse_job_to_file(self, job: Dict):
         """
         Splits the Job to multiple tasks and writes them down in self.local_queue_file.
@@ -435,7 +448,7 @@ class Scheduler(Essential):
 
         logger.debug(f"Testing for chunking {attr} from {job} with skeleton {skeleton}")
         # First of all decide whether we need to chunk current job (or a sub-job if called recursively).
-        if self.needs_chunking(plural(attr), job):
+        if self.needs_chunking(plural(attr), {**job, **skeleton}):
 
             # Force batches to isolate if we shall be dealing with flat data.
             # But we still respect the `max_PARAM_per_batch` if it is provided in job.
@@ -455,7 +468,7 @@ class Scheduler(Essential):
 
                 # This is not the `skeleton` received during the call, but the remaining parts of the `job`,
                 # not related to current `attr`
-                job_skeleton = {k: v for k, v in job.items() if k not in [possible_attr, f"isolate_{plural(attr)}"]}
+                job_skeleton = {k: v for k, v in job.items() if k not in [possible_attr]}
                 logger.debug(f"For {possible_attr} we got current_vals: {current_vals} from {job}, "
                              f"leaving job_skeleton: {job_skeleton}")
 
@@ -558,6 +571,15 @@ class Scheduler(Essential):
                 pass
 
 
+    @staticmethod
+    def get_isolate_attributes_from_job(data: dict) -> Dict:
+        """
+        Get a dictionary with settings for isolation from data.
+        """
+
+        return {k: v for k, v in data.items() if re.match('isolate_[\w]*|max_[\w]*_per_batch', k)}
+
+
     def needs_chunking(self, attr: str, data: Dict) -> bool:
         """
         Recursively analyses the data and identifies if the current level of data should be chunked.
@@ -570,6 +592,8 @@ class Scheduler(Essential):
         attrs = single_or_plural(attr)
         isolate_attrs = [f"isolate_{a}" for a in attrs] + [f"max_{a}_per_batch" for a in attrs]
 
+        root_isolate_attrs = self.get_isolate_attributes_from_job(data)
+
         if any(data[x] for x in isolate_attrs if x in data):
             logger.debug(f"needs_chunking(): Got requirement to isolate {attr} in the current scope: {data}")
             return True
@@ -587,11 +611,9 @@ class Scheduler(Essential):
                 for val in current_vals:
 
                     for name, subdata in val.items():
-                        logger.debug(f"needs_chunking(): Analysing {next_attr} in {subdata}")
-                        if not subdata:
-                            continue
                         logger.debug(f"needs_chunking(): Going recursive for {next_attr} in {subdata}")
-                        if self.needs_chunking(next_attr, subdata):
+                        if isinstance(subdata, dict) and self.needs_chunking(next_attr,
+                                                                             {**subdata, **root_isolate_attrs}):
                             logger.debug(f"needs_chunking(): Returning True for {next_attr} from {subdata}")
                             return True
 


Test patch:
diff --git a/sosw/test/unit/test_scheduler.py b/sosw/test/unit/test_scheduler.py
index 58f127b..bf29e50 100644
--- a/sosw/test/unit/test_scheduler.py
+++ b/sosw/test/unit/test_scheduler.py
@@ -570,6 +570,15 @@ class Scheduler_UnitTestCase(unittest.TestCase):
         self.assertEqual(val['stores']['store_training'], 'some_string')
 
 
+    def test_get_isolate_attributes_from_job(self):
+
+        GOOD = {'isolate_sections': True, 'isolate_Ss': False, 'max_stores_cool_per_batch': 42}
+        BAD = {'sections': True, 'foo': {'baz': 17}}
+
+        result = self.scheduler.get_isolate_attributes_from_job(data={**GOOD, **BAD})
+        self.assertDictEqual(result, GOOD)
+
+
     def test_chunk_job(self):
 
         pl = deepcopy(self.PAYLOAD)
@@ -616,7 +625,7 @@ class Scheduler_UnitTestCase(unittest.TestCase):
         Here we have a tricky case:
 
         `section_weddings` has 3 different `stores`. In `store_music` we have 5 `products`.
-        With max_products_per_batch we should have:
+        For example with max_products_per_batch = 2, we should have:
 
         - store_1
         - store_2
@@ -626,11 +635,17 @@ class Scheduler_UnitTestCase(unittest.TestCase):
         """
         pl = deepcopy(self.PAYLOAD)
         pl['sections']['section_weddings']['stores']['store_music']['max_products_per_batch'] = 2
+        # pl['sections']['section_funerals']['isolate_stores'] = True
+        # pl['isolate_sections'] = True
+        # pl['isolate_stores'] = True
 
         response = self.scheduler.chunk_job(job=pl)
 
         NUMBER_TASKS_EXPECTED = [
             ('sections', 'section_weddings', 5),
+            ('sections', 'section_funerals', 1),
+            ('sections', 'section_conversions', 1),
+            ('sections', 'section_gifts', 1),
         ]
 
         # for row in response:
@@ -643,6 +658,33 @@ class Scheduler_UnitTestCase(unittest.TestCase):
 
         self.assertEqual(batches,
                          list(chunks(pl['sections']['section_weddings']['stores']['store_music']['products'], 2)))
+        # self.assertEqual(1, 2)
+
+
+    def test_chunk_job__root_level_isolate(self):
+        """
+        Tests that `isolate_ATTRs` in the root of the Payload will be respected for chunking in all the nested
+        elements of the Job recursively.
+
+        In the test payload the 'stores' is the second level of nesting attribute. But we pass it in the root of job.
+        """
+        pl = deepcopy(self.PAYLOAD)
+        pl['isolate_stores'] = True
+
+        response = self.scheduler.chunk_job(job=pl)
+
+        NUMBER_TASKS_EXPECTED = [
+            ('sections', 'section_funerals', 2),
+            ('sections', 'section_weddings', 3),
+            ('sections', 'section_conversions', 3),
+            ('sections', 'section_gifts', 1),
+        ]
+
+        # for row in response:
+        #     pprint.pprint(row)
+        #     print('\n')
+
+        self.check_number_of_tasks(NUMBER_TASKS_EXPECTED, response)
 
 
     ### Tests of other methods ###
@@ -863,3 +905,62 @@ class Scheduler_UnitTestCase(unittest.TestCase):
 
         self.scheduler.task_client.dynamo_db_client.get_capacity.return_value = {'read': 10, 'write': 50}
         self.assertEqual(round(self.scheduler._sleeptime_for_dynamo, 2), 0)
+
+
+    def test_apply_job_schema(self):
+        self.scheduler.config['job_schema_variants']['sample_schema_name'] = {
+            'chunkable_attrs': [
+                ('a', {}),
+            ]
+        }
+
+        self.scheduler.parse_job_to_file = MagicMock()
+        self.scheduler.process_file = MagicMock()
+        self.scheduler({'job': {
+            'lambda_name': 'test_lambda',
+            'job_schema_name': 'sample_schema_name'
+            },
+        })
+        self.assertEqual(self.scheduler.config['job_schema']['chunkable_attrs'][0][0], 'a')
+
+
+    def test_apply_job_schema__default(self):
+        self.scheduler.parse_job_to_file = MagicMock()
+        self.scheduler.process_file = MagicMock()
+        self.scheduler({'job': {
+            'lambda_name': 'test_lambda',
+            },
+        })
+
+        self.assertEqual(self.scheduler.config['job_schema']['chunkable_attrs'][0][0], 'b')
+
+
+    def test_apply_job_schema__default_preserved(self):
+        """
+        First test checks a specific job schema name.
+        Second test checks if after calling the scheduler again we overwrite the config and use the default
+        specific job schema.
+
+        """
+
+        self.scheduler.config['job_schema_variants']['sample_schema_name'] = {
+            'chunkable_attrs': [
+                ('a', {}),
+            ]
+        }
+
+        self.scheduler.parse_job_to_file = MagicMock()
+        self.scheduler.process_file = MagicMock()
+        self.scheduler({'job': {
+            'lambda_name': 'test_lambda',
+            'job_schema_name': 'sample_schema_name'
+            },
+        })
+
+        self.assertEqual(self.scheduler.config['job_schema']['chunkable_attrs'][0][0], 'a')
+        self.scheduler({'job': {
+            'lambda_name': 'test_lambda',
+            },
+        })
+
+        self.assertEqual(self.scheduler.config['job_schema']['chunkable_attrs'][0][0], 'b')
diff --git a/sosw/test/variables.py b/sosw/test/variables.py
index 63cabb1..1de067f 100644
--- a/sosw/test/variables.py
+++ b/sosw/test/variables.py
@@ -100,6 +100,13 @@ TEST_SCHEDULER_CONFIG = {
             ('store', {}),
             ('product', {}),
         ]
+    },
+    'job_schema_variants': {
+        'default': {
+            'chunkable_attrs': [
+                ('b', {}),
+            ]
+        }
     }
 }