Implement multiprocessing for the CVM model

src/asim-cvm/configs/settings.yaml

@@ -102,3 +102,37 @@ models:
   - write_cvm_trip_matrices
   - write_tables
 
+multiprocess: True
+fail_fast: True
+num_processes: 12
+
+chunk_size: 0
+chunk_method: hybrid_uss
+chunk_training_mode: disabled
+
+multiprocess_steps:
+  - name: mp_initialize
+    begin: household_attractor
+  - name: mp_accessibility
+    begin: cvm_accessibility
+    slice:
+      tables:
+        - commercial_accessibility
+      exclude: True
+  - name: mp_route_generation
+    begin: route_generation
+    slice:
+      tables:
+        - establishments
+  # route_generation_tnc single-process only
+  # land_use aggregation and duplicate establishment_id under slicing
+  - name: mp_route_generation_tnc
+    begin: route_generation_tnc
+  - name: mp_route
+    begin: route_purpose_and_vehicle
+    slice:
+      tables:
+        - establishments
+        - routes
+  - name: mp_finalize
+    begin: write_cvm_trip_matrices

src/asim-cvm/configs/shadow_pricing.yaml

@@ -0,0 +1,2 @@
+# placeholder required by multiprocessing
+shadow_pricing_models:

src/asim-cvm/extensions/cvm_establishment_attractor.py

@@ -200,3 +200,7 @@ def establishment_attractor(
 
     # write land use table back to state
     state.add_table("land_use", land_use)
+
+    # instantiate commercial_accessibility (must be checkpointed to be used
+    # to slice it)
+    state.get_dataframe("commercial_accessibility")

src/asim-cvm/extensions/cvm_route_generation.py

@@ -32,7 +32,6 @@
 
 _business_type_offset = int(10 ** np.ceil(np.log10(max(BusinessTypes))))
 
-
 @workflow.step
 def route_generation(
     state: State,
@@ -165,6 +164,7 @@ def route_generation(
 
     # sum the two logsums to get the accessibility
     accessibility_df["accessibility"] = accessibility_df['estab_acc_logsum'] + accessibility_df['household_acc_logsum']
+    state.add_table("commercial_accessibility", accessibility_df)
 
     # get the accessibility column for the industry
     join_df = establishments_df.merge(

src/asim-cvm/extensions/cvm_route_stops.py

@@ -549,12 +549,14 @@ def _dwell_time(
             "max_duration"
         ]
 
-        random_dwell_times = scipy.stats.beta.rvs(
+        random_order = state.get_rn_generator().random_for_df(df)
+
+        random_dwell_times = scipy.stats.beta.ppf(
+            random_order[:,0],
             a=alpha,
             b=beta,
             loc=model_settings.min_dwell_time,
             scale=max_duration,
-            size=len(df),
         )
 
         result_list.append(
@@ -622,7 +624,8 @@ def route_stops(
         np.exp(accessibility_df['estab_acc_hh_food']) +
         np.exp(accessibility_df['estab_acc_hh_package'])
     )
-
+    state.add_table("commercial_accessibility", accessibility_df)
+
     route_trip_num = 1
 
     cv_trips = []
@@ -749,8 +752,6 @@ def route_stops(
 
         nonterminated_routes[model_settings.location_settings_estb.RESULT_COL_NAME] = next_stop_location
 
-        np.random.seed(seed=42)
-
         # Choose dwell time
         nonterminated_routes = _dwell_time(
             state,

src/asim-cvm/extensions/cvm_state.py

@@ -1,5 +1,6 @@
 from activitysim.core.workflow import State as workflow_State
 from activitysim.core.workflow.state import StateAttr
+from activitysim.abm.models.util import canonical_ids as cid
 
 from .cvm_settings import CVMSettings
 
@@ -10,3 +11,24 @@ class State(workflow_State):
     """
 
     settings: CVMSettings = StateAttr(CVMSettings)
+
+    # cvm tables that need RNG channel / traceable-table / canonical
+    # index registration on every process start (single or multiprocess),
+    # since checkpoint reloads only tables listed here.
+    _CVM_TABLES = [
+        ("establishments", "establishment_id"),
+        ("establishments_all", "establishment_id_all"),
+        ("routes", "route_id"),
+    ]
+
+    @classmethod
+    def _register_cvm_tables(cls):
+        for table_name, index_name in cls._CVM_TABLES:
+            if table_name not in cid.RANDOM_CHANNELS:
+                cid.RANDOM_CHANNELS.append(table_name)
+            if table_name not in cid.TRACEABLE_TABLES:
+                cid.TRACEABLE_TABLES.append(table_name)
+            cid.CANONICAL_TABLE_INDEX_NAMES.setdefault(table_name, index_name)
+
+
+State._register_cvm_tables()

src/asim-cvm/run_cvm.py

@@ -3,23 +3,40 @@
 from pathlib import Path
 import pytest
 import activitysim.abm  # register components # noqa: F401
+from activitysim.core import mp_tasks
+from activitysim.cli.run import INJECTABLES
 
-# TODO: should this just be its own installable package?
-ext_dir = str(Path(__file__).parents)
-sys.path.insert(0, ext_dir)
-extensions = importlib.import_module("extensions")
-sys.path = sys.path[1:]
+def main():
+    # TODO: should this just be its own installable package?
+    ext_dir = str(Path(__file__).parents)
+    sys.path.insert(0, ext_dir)
+    extensions = importlib.import_module("extensions")
+    sys.path = sys.path[1:]
 
-arguments = sys.argv
+    arguments = sys.argv
 
-data_dir = tuple(arguments[1].split(","))
-configs_dir = tuple(arguments[2].split(","))
-output_dir = arguments[3]
+    data_dir = tuple(arguments[1].split(","))
+    configs_dir = tuple(arguments[2].split(","))
+    output_dir = arguments[3]
 
-state = extensions.cvm_state.State.make_default(
-    __file__, configs_dir=configs_dir, data_dir=data_dir, output_dir=output_dir
-)
+    state = extensions.cvm_state.State.make_default(
+        __file__, configs_dir=configs_dir, data_dir=data_dir, output_dir=output_dir
+    )
 
-state.logging.config_logger()
+    state.import_extensions("extensions")
+    state.logging.config_logger()
 
-state.run(models=state.settings.models)
+    if state.settings.multiprocess:
+        injectables = {}
+        for k in INJECTABLES:
+            try:
+                injectables[k] = state.get_injectable(k)
+            except KeyError:
+                pass
+        injectables["settings"] = state.settings
+        mp_tasks.run_multiprocess(state, injectables)
+    else:
+        state.run(models=state.settings.models)
+
+if __name__ == "__main__":
+    main()