Merge pull request #209 from botify-labs/bugfix/200/try-to-use-less-memory

Bugfix/200/try to use less memory

Author: ybastide

simpleflow/command.py

@@ -26,16 +26,21 @@
 from simpleflow.utils import json_dumps
 from simpleflow import __version__
 
+if False:
+    from typing import Text, Type
+    from simpleflow import Workflow
+    from swf.models import WorkflowType
+
+
 logger = logging.getLogger(__name__)
 
 
 def get_workflow(clspath):
+    # type: (Text) -> Type[Workflow]
     """
     Import a workflow class.
     :param clspath: class path
-    :type clspath: str
     :return:
-    :rtype: simpleflow.workflow.Workflow
     """
     modname, clsname = clspath.rsplit('.', 1)
     module = __import__(modname, fromlist=['*'])
@@ -78,19 +83,17 @@ def cli(ctx, header, format):
     ctx.params['header'] = header
 
 
-def get_workflow_type(domain_name, workflow):
+def get_workflow_type(domain_name, workflow_class):
+    # type: (Text, Type[Workflow]) -> WorkflowType
     """
-    Get or create the given workflow.
+    Get or create the given workflow on SWF.
     :param domain_name:
-    :type domain_name: str
-    :param workflow:
-    :type workflow: simpleflow.workflow.Workflow
+    :param workflow_class:
     :return:
-    :rtype: swf.models.WorkflowType
     """
     domain = swf.models.Domain(domain_name)
     query = swf.querysets.WorkflowTypeQuerySet(domain)
-    return query.get_or_create(workflow.name, workflow.version)
+    return query.get_or_create(workflow_class.name, workflow_class.version)
 
 
 def load_input(input_fp):
@@ -164,24 +167,24 @@ def start_workflow(workflow,
                    input,
                    input_file,
                    local):
-    workflow_definition = get_workflow(workflow)
+    workflow_class = get_workflow(workflow)
 
     wf_input = get_or_load_input(input_file, input)
 
     if local:
         from .local import Executor
 
-        Executor(workflow_definition).run(wf_input)
+        Executor(workflow_class).run(wf_input)
 
         return
 
     if not domain:
         raise ValueError('*domain* must be set when not running in local mode')
 
-    workflow_type = get_workflow_type(domain, workflow_definition)
+    workflow_type = get_workflow_type(domain, workflow_class)
     execution = workflow_type.start_execution(
         workflow_id=workflow_id,
-        task_list=task_list or workflow_definition.task_list,
+        task_list=task_list or workflow_class.task_list,
         execution_timeout=execution_timeout,
         input=wf_input,
         tag_list=tags,

simpleflow/executor.py

@@ -3,6 +3,10 @@
 
 from ._decorators import deprecated
 
+if False:
+    from typing import Type
+    from simpleflow import Workflow
+
 __all__ = ['Executor']
 
 
@@ -27,13 +31,11 @@ class Executor(object):
     - asynchronous
     - asynchronous with full replay
 
-     :ivar _workflow: the workflow
-     :type _workflow: simpleflow.workflow.Workflow
-
     """
     __metaclass__ = abc.ABCMeta
 
-    def __init__(self, workflow):
+    def __init__(self, workflow_class):
+        # type: (Type[Workflow]) -> None
         """
         Binds the workflow's definition.
 
@@ -42,23 +44,30 @@ def __init__(self, workflow):
         as a program, the workflow, and an interpreter, the executor.
 
         """
-        self._workflow = workflow(self)
+        self._workflow_class = workflow_class
+        self._workflow = None
+
+    @property
+    def workflow_class(self):
+        return self._workflow_class
 
     @property
     def workflow(self):
-        """
-        :return:
-        :rtype: simpleflow.workflow.Workflow
-        """
         return self._workflow
 
+    def create_workflow(self):
+        if self._workflow is None:
+            workflow = self._workflow_class(self)
+            if False:
+                assert isinstance(workflow, Workflow)
+            self._workflow = workflow
+
     def run_workflow(self, *args, **kwargs):
         """
         Runs the workflow definition.
 
         """
-        workflow = self._workflow
-        result = workflow.run(*args, **kwargs)
+        result = self._workflow.run(*args, **kwargs)
         return result
 
     @abc.abstractmethod
@@ -78,7 +87,7 @@ def submit(self, submittable, *args, **kwargs):
     def map(self, callable, iterable):
         """Submit *callable* with each of the items in ``*iterables``.
 
-        All items in ``*iterables`` must be serializable in JSON.
+        All items in ``*iterable`` must be serializable in JSON.
 
         """
         return [self.submit(callable, argument) for

simpleflow/local/executor.py

@@ -22,20 +22,20 @@ class Executor(executor.Executor):
     Executes all tasks synchronously in a single local process.
 
     """
-    def __init__(self, workflow):
-        super(Executor, self).__init__(workflow)
+    def __init__(self, workflow_class):
+        super(Executor, self).__init__(workflow_class)
+        self.update_workflow_class()
         self.nb_activities = 0
         self.signals_sent = set()
 
-    @property
-    def _workflow_class(self):
+    def update_workflow_class(self):
         """
         Returns the workflow class with all the needed attributes for
         swf.models.history.builder.History()
         This allows to get a SWF-compatible history in local executions so that
         the metrology feature works correctly.
         """
-        cls = self._workflow.__class__
+        cls = self._workflow_class
         for attr in ("decision_tasks_timeout", "execution_timeout", ):
             if not hasattr(cls, attr):
                 setattr(cls, attr, None)
@@ -105,12 +105,15 @@ def run(self, input=None):
             input = {}
         args = input.get('args', ())
         kwargs = input.get('kwargs', {})
+        self.create_workflow()
 
         self.initialize_history(input)
 
         self.before_replay()
         result = self.run_workflow(*args, **kwargs)
 
+        # Hack: self._history must be available to the callback as a
+        # simpleflow.history.History, not a swf.models.history.builder.History
         self._history = History(self._history)
         self._history.parse()
         self.after_replay()

simpleflow/swf/executor.py

@@ -21,7 +21,6 @@
 from simpleflow.base import Submittable
 from simpleflow.history import History
 from simpleflow.signal import WaitForSignal
-from simpleflow.swf import constants
 from simpleflow.swf.helpers import swf_identity
 from simpleflow.swf.task import ActivityTask, WorkflowTask, SignalTask
 from simpleflow.task import (
@@ -151,7 +150,6 @@ class Executor(executor.Executor):
 
     :ivar domain: domain
     :type domain: swf.models.domain.Domain
-    :ivar workflow: workflow
     :ivar task_list: task list
     :type task_list: Optional[str]
     :ivar repair_with: previous history to use for repairing
@@ -161,9 +159,9 @@ class Executor(executor.Executor):
 
     """
 
-    def __init__(self, domain, workflow, task_list=None, repair_with=None,
+    def __init__(self, domain, workflow_class, task_list=None, repair_with=None,
                  force_activities=None):
-        super(Executor, self).__init__(workflow)
+        super(Executor, self).__init__(workflow_class)
         self._history = None
         self._execution_context = {}
         self.domain = domain
@@ -175,6 +173,7 @@ def __init__(self, domain, workflow, task_list=None, repair_with=None,
             self.force_activities = None
         self.reset()
 
+    # noinspection PyAttributeOutsideInit
     def reset(self):
         """
         Clears the state of the execution.
@@ -187,6 +186,9 @@ def reset(self):
         self._decisions = []
         self._tasks = TaskRegistry()
         self._idempotent_tasks_to_submit = set()
+        self._execution = None
+        self.current_priority = None
+        self.create_workflow()
 
     def _make_task_id(self, a_task, *args, **kwargs):
         """
@@ -783,6 +785,7 @@ def replay(self, decision_response):
                 len(self._decisions),
             ))
             self.after_replay()
+            self.decref_workflow()
             return self._decisions, {}
         except exceptions.TaskException as err:
             reason = 'Workflow execution error in task {}: "{}"'.format(
@@ -799,6 +802,7 @@ def replay(self, decision_response):
                 details=swf.format.details(details),
             )
             self.after_closed()
+            self.decref_workflow()
             return [decision], {}
 
         except Exception as err:
@@ -819,15 +823,23 @@ def replay(self, decision_response):
                 details=swf.format.details(details),
             )
             self.after_closed()
+            self.decref_workflow()
             return [decision], {}
 
         self.after_replay()
         decision = swf.models.decision.WorkflowExecutionDecision()
         decision.complete(result=swf.format.result(json_dumps(result)))
         self.on_completed()
         self.after_closed()
+        self.decref_workflow()
         return [decision], {}
 
+    def decref_workflow(self):
+        """
+        Set the `_workflow` ivar to None in the hope of reducing memory consumption.
+        """
+        self._workflow = None
+
     def before_replay(self):
         return self._workflow.before_replay(self._history)
 

simpleflow/swf/process/decider/base.py

@@ -64,19 +64,19 @@ def __init__(self, workflow_executors, domain, task_list, nb_retries=3,
         """
         self._workflow_name = '{}'.format(','.join(
             [
-                ex.workflow.name for ex in workflow_executors
+                ex.workflow_class.name for ex in workflow_executors
                 ]))
 
         # Maps a workflow's name to its definition.
         # Used to dispatch a decision task to the corresponding workflow.
         self._workflow_executors = {
-            executor.workflow.name: executor for executor in workflow_executors
+            executor.workflow_class.name: executor for executor in workflow_executors
         }
 
         if task_list:
             self.task_list = task_list
         else:
-            self.task_list = workflow_executors[0].workflow.task_list
+            self.task_list = workflow_executors[0].workflow_class.task_list
             # If not passed explicitly, all executors must use the same task list
             # else it's probably a mistake so we raise an error.
             self._check_all_task_lists_identical()
@@ -106,7 +106,7 @@ def _check_all_domains_identical(self):
 
     def _check_all_task_lists_identical(self):
         for ex in self._workflow_executors.values():
-            if ex.workflow.task_list != self.task_list:
+            if ex.workflow_class.task_list != self.task_list:
                 raise ValueError(
                     'all workflows must have the same task list '
                     '"{}" unless you specify it explicitly'.format(

simpleflow/swf/task.py

@@ -32,7 +32,6 @@ def schedule(self, domain, task_list=None, **kwargs):
         :param task_list:
         :type task_list: Optional[str]
         :param kwargs:
-        :type kwargs: dict
         :return:
         :rtype: list[swf.models.decision.Decision]
         """

tests/test_simpleflow/test_dataflow.py

@@ -18,7 +18,6 @@
 from swf.responses import Response
 
 from simpleflow import (
-    Workflow,
     futures,
 )
 from simpleflow.task import ActivityTask
@@ -40,6 +39,9 @@
 )
 
 
+# Note: tests checking the Executor.workflow instance must patch `decref_workflow`.
+
+
 def check_task_scheduled_decision(decision, task):
     """
     Asserts that *decision* schedules *task*.
@@ -218,7 +220,8 @@ def run(self, a):
 
 
 @mock_swf
-def test_workflow_with_before_replay():
+@patch.object(Executor, 'decref_workflow')
+def test_workflow_with_before_replay(mock_decref_workflow):
     workflow = ATestDefinitionWithBeforeReplay
     executor = Executor(DOMAIN, workflow)
 
@@ -248,7 +251,8 @@ def run(self, a):
 
 
 @mock_swf
-def test_workflow_with_after_replay():
+@patch.object(Executor, 'decref_workflow')
+def test_workflow_with_after_replay(mock_decref_workflow):
     workflow = ATestDefinitionWithAfterReplay
     executor = Executor(DOMAIN, workflow)
 
@@ -277,7 +281,8 @@ def run(self, a):
 
 
 @mock_swf
-def test_workflow_with_after_closed():
+@patch.object(Executor, 'decref_workflow')
+def test_workflow_with_after_closed(mock_decref_workflow):
     workflow = ATestDefinitionWithAfterClosed
     executor = Executor(DOMAIN, workflow)
 
@@ -1068,7 +1073,8 @@ def run(self):
 
 
 @mock_swf
-def test_workflow_failed_from_definition():
+@patch.object(Executor, 'decref_workflow')
+def test_workflow_failed_from_definition(mock_decref_workflow):
     workflow = ATestDefinitionFailWorkflow
     executor = Executor(DOMAIN, workflow)
     history = builder.History(workflow)
@@ -1110,7 +1116,8 @@ def run(self):
 
 
 @mock_swf
-def test_workflow_activity_raises_on_failure():
+@patch.object(Executor, 'decref_workflow')
+def test_workflow_activity_raises_on_failure(mock_decref_workflow):
     workflow = ATestDefinitionActivityRaisesOnFailure
     executor = Executor(DOMAIN, workflow)
     history = builder.History(workflow)
@@ -1148,7 +1155,8 @@ def run(self):
 
 
 @mock_swf
-def test_on_failure_callback():
+@patch.object(Executor, 'decref_workflow')
+def test_on_failure_callback(mock_decref_workflow):
     workflow = ATestOnFailureDefinition
     executor = Executor(DOMAIN, workflow)
     history = builder.History(workflow)