diff --git a/activitysim/abm/models/location_choice.py b/activitysim/abm/models/location_choice.py index c955f4a606..662f7bf264 100644 --- a/activitysim/abm/models/location_choice.py +++ b/activitysim/abm/models/location_choice.py @@ -11,7 +11,15 @@ from activitysim.abm.models.util import tour_destination from activitysim.abm.models.util.bias_logsums import maybe_bias_logsums from activitysim.abm.tables import shadow_pricing -from activitysim.core import estimation, expressions, los, simulate, tracing, workflow +from activitysim.core import ( + chunk, + estimation, + expressions, + los, + simulate, + tracing, + workflow, +) from activitysim.core.configuration.logit import ( TourLocationComponentSettings, TourModeComponentSettings, @@ -520,8 +528,7 @@ def run_location_sample( full_dest_size_terms = dest_size_terms logger.debug( - f"dropping {(~(dest_size_terms.size_term > 0)).sum()} " - f"of {len(dest_size_terms)} rows where size_term is zero" + f"dropping {(~(dest_size_terms.size_term > 0)).sum()} of {len(dest_size_terms)} rows where size_term is zero" ) dest_size_terms = dest_size_terms[dest_size_terms.size_term > 0] @@ -532,8 +539,7 @@ def run_location_sample( if pre_sample_taz and not state.settings.want_dest_choice_presampling: pre_sample_taz = False logger.info( - f"Disabled destination zone presampling for {trace_label} " - f"because 'want_dest_choice_presampling' setting is False" + f"Disabled destination zone presampling for {trace_label} because 'want_dest_choice_presampling' setting is False" ) if pre_sample_taz: @@ -616,23 +622,56 @@ def run_location_logsums( logger.info(f"Running {trace_label} with {len(location_sample_df.index)} rows") - choosers = location_sample_df.join(persons_merged_df, how="left") - tour_purpose = model_settings.LOGSUM_TOUR_PURPOSE if isinstance(tour_purpose, dict): tour_purpose = tour_purpose[segment_name] - logsums = logsum.compute_location_choice_logsums( + # Join sampled alternatives to person attributes inside the chooser chunk. + # The old full-table join could retain millions of rows and all derived + # logsum preprocessor columns before the utility evaluator began chunking. + # At production scale that defeated explicit chunking and exhausted memory. + pnr_index_multiplier = logsum.get_pnr_index_multiplier( + location_sample_df, logsum_settings + ) + logsum_chunks = [] + for ( + _i, + persons_chunk, + location_sample_chunk, + chunk_trace_label, + chunk_sizer, + ) in chunk.adaptive_chunked_choosers_and_alts( state, - choosers, - tour_purpose, - logsum_settings, - model_settings, - network_los, - chunk_size, - chunk_tag, + persons_merged_df, + location_sample_df, trace_label, - ) + chunk_tag, + chunk_size=chunk_size, + explicit_chunk_size=model_settings.explicit_chunk, + ): + choosers = location_sample_chunk.join(persons_chunk, how="left") + assert choosers.index.equals(location_sample_chunk.index) + chunk_sizer.log_df(chunk_trace_label, "logsum_choosers", choosers) + + logsum_chunks.append( + logsum.compute_location_choice_logsums( + state, + choosers, + tour_purpose, + logsum_settings, + model_settings, + network_los, + 0, + chunk_tag, + chunk_trace_label, + explicit_chunk_size=0, + pnr_index_multiplier=pnr_index_multiplier, + ) + ) + chunk_sizer.log_df(chunk_trace_label, "logsum_choosers", None) + + logsums = pd.concat(logsum_chunks) + assert logsums.index.equals(location_sample_df.index) # "add_column series should have an index matching the table to which it is being added" # when the index has duplicates, however, in the special case that the series index exactly diff --git a/activitysim/abm/models/park_and_ride_lot_choice.py b/activitysim/abm/models/park_and_ride_lot_choice.py index 4cebc4020a..d99895a388 100644 --- a/activitysim/abm/models/park_and_ride_lot_choice.py +++ b/activitysim/abm/models/park_and_ride_lot_choice.py @@ -182,6 +182,9 @@ def run_park_and_ride_lot_choice( model_settings_file_name: str = "park_and_ride_lot_choice.yaml", pnr_capacity_cls: ParkAndRideCapacity | None = None, trace_label: str = "park_and_ride_lot_choice", + chunk_size: int | None = None, + explicit_chunk_size: float | None = None, + chooser_index_multiplier: int | None = None, ) -> pd.Series: """ Run the park-and-ride lot choice model. @@ -217,7 +220,7 @@ def run_park_and_ride_lot_choice( pnr_alts["pnr_lot_full"] = 0 original_index = None - if not choosers.index.is_unique: + if chooser_index_multiplier is not None or not choosers.index.is_unique: # non-unique index will crash interaction_simulate # so we need to reset the index and add it to ActivitySim's rng # this happens while the disaggregate accessibility model is running pnr lot choice @@ -225,9 +228,13 @@ def run_park_and_ride_lot_choice( oi_name = original_index.name oi_name = oi_name if oi_name else "index" choosers = choosers.reset_index(drop=False) - idx_multiplier = choosers.groupby(oi_name).size().max() - # round to the nearest 10's place - idx_multiplier = int(np.ceil(idx_multiplier / 10.0) * 10) + # A logsum caller supplies the full segment's multiplier. Use it even + # when this chunk happens to have unique indices, so all chunks use the + # same synthetic RNG channel and keys as the unchunked segment. + idx_multiplier = chooser_index_multiplier + if idx_multiplier is None: + idx_multiplier = choosers.groupby(oi_name).size().max() + idx_multiplier = int(np.ceil(idx_multiplier / 10.0) * 10) choosers.index = ( original_index * idx_multiplier + choosers.groupby(oi_name).cumcount() ) @@ -345,7 +352,12 @@ def run_park_and_ride_lot_choice( trace_label=trace_label, trace_choice_name=trace_label, estimator=estimator, - explicit_chunk_size=model_settings.explicit_chunk, + chunk_size=chunk_size, + explicit_chunk_size=( + model_settings.explicit_chunk + if explicit_chunk_size is None + else explicit_chunk_size + ), compute_settings=model_settings.compute_settings, ) diff --git a/activitysim/abm/models/trip_destination.py b/activitysim/abm/models/trip_destination.py index 7dc6154cc5..5f8157a0a1 100644 --- a/activitysim/abm/models/trip_destination.py +++ b/activitysim/abm/models/trip_destination.py @@ -26,11 +26,11 @@ estimation, expressions, los, + mem, simulate, tracing, workflow, ) -from activitysim.core.configuration.base import PreprocessorSettings from activitysim.core.configuration.logit import LocationComponentSettings from activitysim.core.exceptions import DuplicateWorkflowTableError, InvalidTravelError from activitysim.core.interaction_sample import ( @@ -82,8 +82,7 @@ def deprecated_destination_prefix(cls, values): if values[badkey] != values[goodkey]: # both keys are given, with different values -> error raise ValueError( - f"Deprecated `{badkey}` field must have the " - f"same value as `{goodkey}` if both are provided." + f"Deprecated `{badkey}` field must have the same value as `{goodkey}` if both are provided." ) else: # both keys are given, with same values -> warning @@ -98,9 +97,7 @@ def deprecated_destination_prefix(cls, values): else: # only the wrong key is given -> warning warnings.warn( - f"Use of the field `{badkey}` in the " - "trip_destination configuration file is deprecated, use " - f"`{goodkey}` instead.", + f"Use of the field `{badkey}` in the trip_destination configuration file is deprecated, use `{goodkey}` instead.", FutureWarning, stacklevel=2, ) @@ -149,6 +146,7 @@ def _destination_sample( chunk_tag: str, trace_label: str, zone_layer=None, + preprocess_alternatives: bool = True, ): """ @@ -209,16 +207,17 @@ def _destination_sample( log_alt_losers = state.settings.log_alt_losers - # preprocessing alternatives - expressions.annotate_preprocessors( - state, - df=alternatives, - locals_dict=locals_dict, - skims=skims, - model_settings=model_settings, - trace_label=trace_label, - preprocessor_setting_name="alts_preprocessor_sample", - ) + if preprocess_alternatives: + # preprocessing alternatives + expressions.annotate_preprocessors( + state, + df=alternatives, + locals_dict=locals_dict, + skims=skims, + model_settings=model_settings, + trace_label=trace_label, + preprocessor_setting_name="alts_preprocessor_sample", + ) # Trip destination keeps the alternative universe here so stable_alt_positions is not needed. choices = interaction_sample( @@ -694,33 +693,78 @@ def destination_presample( skims = skim_hotel.sample_skims(presample=True) - taz_sample = _destination_sample( - state, - primary_purpose, - trips_taz, - alternatives, - model_settings, - TAZ_size_term_matrix, - skims, - alt_dest_col_name, - estimator, - chunk_tag=chunk_tag, - trace_label=trace_label, - zone_layer="taz", - ) + explicit_chunk_size = getattr(model_settings, "explicit_chunk", 0) + if ( + explicit_chunk_size + and state.settings.chunk_training_mode != chunk.MODE_EXPLICIT + ): + # Adaptive sampling already owns a ledger; do not nest an additional + # pipeline ledger merely because the model has an explicit setting. + explicit_chunk_size = 0 + if explicit_chunk_size: + chooser_chunks = chunk.adaptive_chunked_choosers( + state, + trips_taz, + trace_label, + f"{chunk_tag}.pipeline", + chunk_size=state.settings.chunk_size, + explicit_chunk_size=explicit_chunk_size, + ) + else: + # Preserve the legacy unchunked call contract, including compatibility + # with callers that supply a lightweight settings object. + chooser_chunks = ((0, trips_taz, trace_label, None),) + + # Bound the entire two-stage sampling pipeline, not just utility + # evaluation inside interaction_sample. choose_MAZ_for_TAZ temporarily + # expands every sampled TAZ by its constituent MAZs; doing that for a full + # trip-purpose segment can consume many GiB even when interaction_sample is + # itself chunked. Keeping both stages inside this outer chooser loop lets + # each expanded MAZ frame be released before the next chunk. + maz_sample_chunks = [] + preprocess_alternatives = True + for _i, trips_taz_chunk, chunk_trace_label, _chunk_sizer in chooser_chunks: + sample_kwargs = {} + if explicit_chunk_size: + sample_kwargs["preprocess_alternatives"] = preprocess_alternatives + + taz_sample = _destination_sample( + state, + primary_purpose, + trips_taz_chunk, + alternatives, + model_settings, + TAZ_size_term_matrix, + skims, + alt_dest_col_name, + estimator, + chunk_tag=chunk_tag, + trace_label=chunk_trace_label, + zone_layer="taz", + **sample_kwargs, + ) + if explicit_chunk_size: + preprocess_alternatives = False - # choose a MAZ for each DEST_TAZ choice, choice probability based on MAZ size_term fraction of TAZ total - maz_sample = choose_MAZ_for_TAZ( - state, - taz_sample, - size_term_matrix, - trips, - network_los, - alt_dest_col_name, - trace_label, - model_settings, - full_taz_index=full_taz_index, - ) + # Choose a MAZ for each DEST_TAZ choice, with probability based on its + # share of the TAZ's purpose-specific size term. + maz_sample_chunk = choose_MAZ_for_TAZ( + state, + taz_sample, + size_term_matrix, + trips, + network_los, + alt_dest_col_name, + chunk_trace_label, + model_settings, + full_taz_index=full_taz_index, + ) + maz_sample_chunks.append(maz_sample_chunk) + del taz_sample, maz_sample_chunk + if explicit_chunk_size: + mem.release_memory() + + maz_sample = pd.concat(maz_sample_chunks) assert alt_dest_col_name in maz_sample @@ -766,8 +810,7 @@ def trip_destination_sample( if pre_sample_taz and not state.settings.want_dest_choice_presampling: pre_sample_taz = False logger.info( - f"Disabled destination zone presampling for {trace_label} " - f"because 'want_dest_choice_presampling' setting is False" + f"Disabled destination zone presampling for {trace_label} because 'want_dest_choice_presampling' setting is False" ) if pre_sample_taz: @@ -878,27 +921,12 @@ def compute_logsums( # chunk usage is uniform so better to combine chunk_tag = "trip_destination.compute_logsums" - # FIXME should pass this in? - network_los = state.get_injectable("network_los") - # - trips_merged - merge trips and tours_merged trips_merged = pd.merge( trips, tours_merged, left_on="tour_id", right_index=True, how="left" ) assert trips_merged.index.equals(trips.index) - # - choosers - merge destination_sample and trips_merged - # re/set index because pandas merge does not preserve left index if it has duplicate values! - choosers = pd.merge( - destination_sample, - trips_merged.reset_index(), - left_index=True, - right_on="trip_id", - how="left", - suffixes=("", "_r"), - ).set_index("trip_id") - assert choosers.index.equals(destination_sample.index) - logsum_settings = state.filesystem.read_model_settings( model_settings.LOGSUM_SETTINGS ) @@ -932,20 +960,6 @@ def compute_logsums( "timeframe": "trip", } - destination_sample["od_logsum"] = compute_ood_logsums( - state, - choosers, - logsum_settings, - nest_spec, - logsum_spec, - od_skims, - locals_dict, - state.settings.chunk_size, - trace_label=tracing.extend_trace_label(trace_label, "od"), - chunk_tag=chunk_tag, - explicit_chunk_size=model_settings.explicit_chunk, - ) - # - dp_logsums dp_skims = { "ORIGIN": model_settings.ALT_DEST_COL_NAME, @@ -955,19 +969,83 @@ def compute_logsums( "od_skims": skims["dp_skims"], } - destination_sample["dp_logsum"] = compute_ood_logsums( + # Merge sampled alternatives with chooser attributes inside the chunk loop. + # Previously this merge and both preprocessors ran on the full sample table + # before ``simple_simulate_logsums`` chunked utility evaluation. For large + # trip models, that briefly retained millions of rows and dozens of derived + # columns per process, defeating the purpose of explicit chunking. + od_logsum_chunks = [] + dp_logsum_chunks = [] + for ( + _i, + trips_chunk, + destination_sample_chunk, + chunk_trace_label, + chunk_sizer, + ) in chunk.adaptive_chunked_choosers_and_alts( state, - choosers, - logsum_settings, - nest_spec, - logsum_spec, - dp_skims, - locals_dict, - state.settings.chunk_size, - trace_label=tracing.extend_trace_label(trace_label, "dp"), - chunk_tag=chunk_tag, + trips_merged, + destination_sample, + trace_label, + chunk_tag, + chunk_size=state.settings.chunk_size, explicit_chunk_size=model_settings.explicit_chunk, - ) + ): + # Re/set the index because pandas merge does not preserve the left index + # when it contains the repeated trip ids of sampled alternatives. + choosers = pd.merge( + destination_sample_chunk, + trips_chunk.reset_index(), + left_index=True, + right_on="trip_id", + how="left", + suffixes=("", "_r"), + ).set_index("trip_id") + assert choosers.index.equals(destination_sample_chunk.index) + chunk_sizer.log_df(chunk_trace_label, "logsum_choosers", choosers) + + # The outer loop now owns chunking. Disable nested chunking so each + # merged/preprocessed sample chunk can be released before constructing + # the next one. + od_logsum_chunk = compute_ood_logsums( + state, + choosers, + logsum_settings, + nest_spec, + logsum_spec, + od_skims, + locals_dict, + 0, + trace_label=tracing.extend_trace_label(chunk_trace_label, "od"), + chunk_tag=chunk_tag, + explicit_chunk_size=0, + ) + dp_logsum_chunk = compute_ood_logsums( + state, + choosers, + logsum_settings, + nest_spec, + logsum_spec, + dp_skims, + locals_dict, + 0, + trace_label=tracing.extend_trace_label(chunk_trace_label, "dp"), + chunk_tag=chunk_tag, + explicit_chunk_size=0, + ) + od_logsum_chunks.append(od_logsum_chunk) + dp_logsum_chunks.append(dp_logsum_chunk) + chunk_sizer.log_df(chunk_trace_label, "logsum_choosers", None) + del choosers, od_logsum_chunk, dp_logsum_chunk + mem.release_memory() + + od_logsums = pd.concat(od_logsum_chunks) + dp_logsums = pd.concat(dp_logsum_chunks) + assert od_logsums.index.equals(destination_sample.index) + assert dp_logsums.index.equals(destination_sample.index) + + destination_sample["od_logsum"] = od_logsums + destination_sample["dp_logsum"] = dp_logsums return destination_sample @@ -1026,9 +1104,6 @@ def trip_destination_simulate( trip_period_idx = skims["odt_skims"].map_time_periods(trips) if trip_period_idx is not None: trips["trip_period"] = trip_period_idx - else: - None - locals_dict = model_settings.CONSTANTS.copy() locals_dict.update( { @@ -1040,38 +1115,94 @@ def trip_destination_simulate( ) locals_dict.update(skims) - # preprocessing alternatives - expressions.annotate_preprocessors( - state, - df=destination_sample, - locals_dict=locals_dict, - skims=skims, - model_settings=model_settings, - trace_label=trace_label, - preprocessor_setting_name="alts_preprocessor_simulate", - ) - log_alt_losers = state.settings.log_alt_losers - destinations = interaction_sample_simulate( - state, - choosers=trips, - alternatives=destination_sample, - spec=spec, - choice_column=alt_dest_col_name, - log_alt_losers=log_alt_losers, - want_logsums=want_logsums, - allow_zero_probs=True, - zero_prob_choice_val=NO_DESTINATION, - skims=skims, - locals_d=locals_dict, - chunk_size=state.settings.chunk_size, - chunk_tag=chunk_tag, - trace_label=trace_label, - trace_choice_name="trip_dest", - estimator=estimator, - explicit_chunk_size=model_settings.explicit_chunk, - alts_context=alts_context, - ) + if estimator: + # Preserve the estimator's single-call lifecycle and output bundle. + expressions.annotate_preprocessors( + state, + df=destination_sample, + locals_dict=locals_dict, + skims=skims, + model_settings=model_settings, + trace_label=trace_label, + preprocessor_setting_name="alts_preprocessor_simulate", + ) + destinations = interaction_sample_simulate( + state, + choosers=trips, + alternatives=destination_sample, + spec=spec, + choice_column=alt_dest_col_name, + log_alt_losers=log_alt_losers, + want_logsums=want_logsums, + allow_zero_probs=True, + zero_prob_choice_val=NO_DESTINATION, + skims=skims, + locals_d=locals_dict, + chunk_size=state.settings.chunk_size, + chunk_tag=chunk_tag, + trace_label=trace_label, + trace_choice_name="trip_dest", + estimator=estimator, + explicit_chunk_size=model_settings.explicit_chunk, + alts_context=alts_context, + ) + else: + # The alternative preprocessor can add several derived columns. Running + # it on the complete sampled table before interaction_sample_simulate + # briefly materializes millions of rows per worker and defeats the + # evaluator's internal chunking. Preprocess and simulate each sampled + # chooser chunk end-to-end instead. + destination_chunks = [] + for ( + _i, + trips_chunk, + destination_sample_chunk, + chunk_trace_label, + _chunk_sizer, + ) in chunk.adaptive_chunked_choosers_and_alts( + state, + trips, + destination_sample, + trace_label, + chunk_tag, + chunk_size=state.settings.chunk_size, + explicit_chunk_size=model_settings.explicit_chunk, + ): + expressions.annotate_preprocessors( + state, + df=destination_sample_chunk, + locals_dict=locals_dict, + skims=skims, + model_settings=model_settings, + trace_label=chunk_trace_label, + preprocessor_setting_name="alts_preprocessor_simulate", + ) + destination_chunk = interaction_sample_simulate( + state, + choosers=trips_chunk, + alternatives=destination_sample_chunk, + spec=spec, + choice_column=alt_dest_col_name, + log_alt_losers=log_alt_losers, + want_logsums=want_logsums, + allow_zero_probs=True, + zero_prob_choice_val=NO_DESTINATION, + skims=skims, + locals_d=locals_dict, + chunk_size=0, + chunk_tag=chunk_tag, + trace_label=chunk_trace_label, + trace_choice_name="trip_dest", + estimator=None, + explicit_chunk_size=0, + alts_context=alts_context, + ) + destination_chunks.append(destination_chunk) + del destination_chunk + mem.release_memory() + + destinations = pd.concat(destination_chunks) if not want_logsums: # for consistency, always return a dataframe with canonical column name @@ -1095,8 +1226,7 @@ def trip_destination_simulate( return destinations -@workflow.func -def choose_trip_destination( +def _choose_trip_destination_unchunked( state: workflow.State, primary_purpose, trips, @@ -1136,8 +1266,7 @@ def choose_trip_destination( dropped_trips = ~trips.index.isin(destination_sample.index.unique()) if dropped_trips.any(): logger.warning( - "%s trip_destination_sample %s trips " - "without viable destination alternatives" + "%s trip_destination_sample %s trips without viable destination alternatives" % (trace_label, dropped_trips.sum()) ) trips = trips[~dropped_trips] @@ -1186,8 +1315,7 @@ def choose_trip_destination( dropped_trips = ~trips.index.isin(destinations.index) if dropped_trips.any(): logger.warning( - "%s trip_destination_simulate %s trips " - "without viable destination alternatives" + "%s trip_destination_simulate %s trips without viable destination alternatives" % (trace_label, dropped_trips.sum()) ) @@ -1206,6 +1334,98 @@ def choose_trip_destination( return destinations, destination_sample +@workflow.func +def choose_trip_destination( + state: workflow.State, + primary_purpose, + trips, + alternatives, + tours_merged, + model_settings: TripDestinationSettings, + want_logsums, + want_sample_table, + size_term_matrix, + skim_hotel, + estimator, + chunk_size, + trace_label, +): + """Run the complete destination-choice pipeline in bounded chooser chunks. + + Sampling, logsum calculation, and final simulation used to be chunked + independently. That bounded each temporary calculation, but the complete + sampled-alternative table for a purpose remained live between stages. For + large trip-purpose segments this retained millions of rows per worker and + made memory grow again as the next trip number began. + + When explicit chunking is configured, stream each chooser chunk through + all three stages and retain only its final one-row-per-trip choice. The + estimator and sample-table paths keep their original whole-segment + lifecycle because they intentionally consume the complete sample table. + """ + + if ( + state.settings.chunk_training_mode != chunk.MODE_EXPLICIT + or estimator + or want_sample_table + or not model_settings.explicit_chunk + ): + return _choose_trip_destination_unchunked( + state, + primary_purpose, + trips, + alternatives, + tours_merged, + model_settings, + want_logsums, + want_sample_table, + size_term_matrix, + skim_hotel, + estimator, + chunk_size, + trace_label, + ) + + # The outer pipeline owns the chunk boundary. In particular, a fractional + # size must not be applied again to each already bounded chooser chunk. + inner_settings = model_settings.model_copy(update={"explicit_chunk": 0}) + destination_chunks = [] + for ( + _i, + trips_chunk, + chunk_trace_label, + _chunk_sizer, + ) in chunk.adaptive_chunked_choosers( + state, + trips, + trace_label, + "trip_destination.pipeline", + chunk_size=state.settings.chunk_size, + explicit_chunk_size=model_settings.explicit_chunk, + ): + destinations_chunk, destination_sample = _choose_trip_destination_unchunked( + state, + primary_purpose, + trips_chunk, + alternatives, + tours_merged, + inner_settings, + want_logsums, + want_sample_table=False, + size_term_matrix=size_term_matrix, + skim_hotel=skim_hotel, + estimator=None, + chunk_size=chunk_size, + trace_label=chunk_trace_label, + ) + assert destination_sample is None + destination_chunks.append(destinations_chunk) + del destinations_chunk, destination_sample + mem.release_memory() + + return pd.concat(destination_chunks), None + + class SkimHotel: def __init__( self, @@ -1331,8 +1551,6 @@ def run_trip_destination( model_settings = TripDestinationSettings.read_settings_file( state.filesystem, model_settings_file_name ) - preprocessor_settings = model_settings.preprocessor - # read in logsum settings if they exist, otherwise logsum calculations are skipped if model_settings.LOGSUM_SETTINGS: logsum_settings = state.filesystem.read_model_settings( @@ -1481,9 +1699,6 @@ def run_trip_destination( ) if trip_period_idx is not None: nth_trips["trip_period"] = trip_period_idx - else: - None - logger.debug( "Running %s with %d trips", nth_trace_label, nth_trips.shape[0] ) @@ -1535,6 +1750,14 @@ def run_trip_destination( assert destination_sample is not None sample_list.append(destination_sample) + # Each purpose can create very large temporary TAZ-to-MAZ and + # logsum arrays. Explicit chunking bounds one allocation, while + # allocator pressure relief prevents freed chunks from + # accumulating in a long-lived multiprocess worker before the + # next purpose begins. + if model_settings.explicit_chunk: + mem.release_memory() + destinations_df = pd.concat(choices_list) if fail_some_trips_for_testing: diff --git a/activitysim/abm/models/trip_mode_choice.py b/activitysim/abm/models/trip_mode_choice.py index 4aad9dda62..9ef9d7cfd7 100644 --- a/activitysim/abm/models/trip_mode_choice.py +++ b/activitysim/abm/models/trip_mode_choice.py @@ -3,24 +3,21 @@ from __future__ import annotations import logging -from typing import Any -import numpy as np import pandas as pd from activitysim.abm.models.util import school_escort_tours_trips from activitysim.abm.models.util.mode import mode_choice_simulate from activitysim.core import ( - chunk, config, estimation, expressions, los, + mem, simulate, tracing, workflow, ) -from activitysim.core.configuration.base import PreprocessorSettings, PydanticReadable from activitysim.core.configuration.logit import TemplatedLogitComponentSettings from activitysim.core.util import assign_in_place @@ -44,7 +41,7 @@ class TripModeChoiceSettings(TemplatedLogitComponentSettings, extra="forbid"): FORCE_ESCORTEE_CHAUFFEUR_MODE_MATCH: bool = True """ If True, overwrite the trip mode of escortee trips to match the mode selected - by the chauffeur. This is useful for school escort tours where the escortee trip + by the chauffeur. This is useful for school escort tours where the escortee trip mode (e.g., "transit") should match the chauffeur trip mode. """ @@ -96,21 +93,12 @@ def trip_mode_choice( else: tours_merged = pd.DataFrame() - # - trips_merged - merge trips and tours_merged - trips_merged = pd.merge( - trips_df, tours_merged, left_on="tour_id", right_index=True, how="left" - ) - assert trips_merged.index.equals(trips.index) - tracing.print_summary( "primary_purpose", trips_df.primary_purpose, value_counts=True ) # setup skim keys - assert "trip_period" not in trips_merged - trips_merged["trip_period"] = network_los.skim_time_period_label( - trips_merged.depart, as_cat=True - ) + assert "trip_period" not in trips_df orig_col = "origin" dest_col = "destination" @@ -130,6 +118,17 @@ def trip_mode_choice( skim_dict = network_los.get_default_skim_dict() + def add_trip_period(choosers): + choosers["trip_period"] = network_los.skim_time_period_label( + choosers.depart, as_cat=True + ) + if hasattr(skim_dict, "map_time_periods_from_series"): + trip_period_idx = skim_dict.map_time_periods_from_series( + choosers["trip_period"] + ) + if trip_period_idx is not None: + choosers["trip_period"] = trip_period_idx + odt_skim_stack_wrapper = skim_dict.wrap_3d( orig_key=orig_col, dest_key=dest_col, dim3_key="trip_period" ) @@ -138,13 +137,6 @@ def trip_mode_choice( ) od_skim_wrapper = skim_dict.wrap("origin", "destination") - if hasattr(skim_dict, "map_time_periods_from_series"): - trip_period_idx = skim_dict.map_time_periods_from_series( - trips_merged["trip_period"] - ) - if trip_period_idx is not None: - trips_merged["trip_period"] = trip_period_idx - skims = { "odt_skims": odt_skim_stack_wrapper, "dot_skims": dot_skim_stack_wrapper, @@ -169,11 +161,27 @@ def trip_mode_choice( choices_list = [] cols_to_keep_list = [] - for primary_purpose, trips_segment in trips_merged.groupby( + for primary_purpose, base_trips_segment in trips_df.groupby( "primary_purpose", observed=True ): segment_trace_label = tracing.extend_trace_label(trace_label, primary_purpose) + # A full trips/tours merge duplicates several gigabytes at production + # scale. Materialize only the chooser rows for the active purpose. + if len(tours_cols) > 0: + trips_segment = pd.merge( + base_trips_segment, + tours_merged, + left_on="tour_id", + right_index=True, + how="left", + ) + else: + trips_segment = base_trips_segment.copy() + assert trips_segment.index.equals(base_trips_segment.index) + + add_trip_period(trips_segment) + logger.info( "trip_mode_choice tour_type '%s' (%s trips)" % ( @@ -266,6 +274,14 @@ def trip_mode_choice( ), "{cols_not_in_choosers} from CHOOSER_COLS_TO_KEEP is not in the choosers dataframe" cols_to_keep_list.append(trips_segment[cols_to_keep]) + # Wrappers retain their last chooser frame (and possibly array views). + # Retarget to an independent empty frame before releasing this purpose. + simulate.set_skim_wrapper_targets( + trips_segment.iloc[:0].copy(), skims, allow_partial_success=False + ) + del trips_segment, base_trips_segment + mem.release_memory() + choices_df = pd.concat(choices_list) if estimator: @@ -295,8 +311,6 @@ def trip_mode_choice( ) ) - tracing.print_summary("trip_modes", trips_merged.tour_mode, value_counts=True) - tracing.print_summary( "trip_mode_choice choices", trips_df[mode_column_name], value_counts=True ) @@ -330,15 +344,28 @@ def trip_mode_choice( # need to update locals_dict to access skims that are the same .shape as trips table locals_dict = {} locals_dict.update(constants) - if state.settings.skip_failed_choices: - trips_merged = trips_merged.loc[~mask_skipped] - simulate.set_skim_wrapper_targets(trips_merged, skims) locals_dict.update(skims) locals_dict["timeframe"] = "trip" - expressions.annotate_tables( - state, - locals_dict=locals_dict, - skims=skims, - model_settings=model_settings, - trace_label=trace_label, - ) + # Three-dimensional skim wrappers require trip_period. It is normally only + # needed in the purpose-sized chooser frames above, but post-choice table + # annotators may also use those skims. Add it to the full trips table only + # for annotation, then restore the original table schema. + temporary_trip_period = "trip_period" not in trips_df.columns + if temporary_trip_period: + add_trip_period(trips_df) + try: + expressions.annotate_tables( + state, + locals_dict=locals_dict, + skims=skims, + model_settings=model_settings, + trace_label=trace_label, + ) + finally: + # CHOOSER_COLS_TO_KEEP may have made trip_period an output column. + # Remove it only when this annotation block created it temporarily. + if temporary_trip_period: + trips_df.drop(columns="trip_period", inplace=True) + state_trips = state.get_dataframe("trips", as_copy=False) + if state_trips is not trips_df: + state_trips.drop(columns="trip_period", inplace=True) diff --git a/activitysim/abm/models/trip_scheduling.py b/activitysim/abm/models/trip_scheduling.py index 0ce4f2ffff..da4a595770 100644 --- a/activitysim/abm/models/trip_scheduling.py +++ b/activitysim/abm/models/trip_scheduling.py @@ -15,10 +15,18 @@ split_out_school_escorting_trips, ) from activitysim.abm.models.util.trip import cleanup_failed_trips, failed_trip_cohorts -from activitysim.core import chunk, config, estimation, expressions, tracing, workflow +from activitysim.core import ( + chunk, + config, + estimation, + expressions, + mem, + tracing, + workflow, +) from activitysim.core.configuration.base import PreprocessorSettings, PydanticReadable -from activitysim.core.util import reindex from activitysim.core.exceptions import InvalidTravelError, PipelineError +from activitysim.core.util import reindex logger = logging.getLogger(__name__) @@ -521,6 +529,7 @@ def trip_scheduling( ) trips_df = trips.copy() + original_trip_columns = trips.columns if state.is_table("school_escort_trips"): school_escort_trips = state.get_dataframe("school_escort_trips") @@ -630,14 +639,36 @@ def trip_scheduling( choices_list.append(choices) - trips_df = trips.copy() - if state.is_table("school_escort_trips"): + # The working frame contains only non-school-escort trips, so rebuild + # the complete output frame as before. Release it first to avoid + # briefly retaining two full trip-table copies. + del trips_chunk, trips_df + mem.release_memory() + trips_df = trips.copy() + # separate out school escorting trips to exclude them from the model and estimation data bundle trips_df, se_trips_df, full_trips_index = split_out_school_escorting_trips( trips_df, school_escort_trips ) non_se_trips_df = trips_df + else: + # ``trips_df`` began as an exact copy of ``trips`` and scheduling only + # added working columns to this outer frame. Reuse it instead of + # allocating a second complete copy of a potentially 30-million-row + # trip table merely to discard those columns. + temporary_columns = trips_df.columns.difference(original_trip_columns) + if len(temporary_columns): + trips_df.drop(columns=temporary_columns, inplace=True) + + # Preserve the old reset-to-``trips`` semantics for the unlikely case + # where a caller supplied columns that trip scheduling also uses as + # scratch space. + overwritten_columns = original_trip_columns.intersection( + ["earliest", "latest", "tour_hour", "stop_num", "chunk_id"] + ) + for column in overwritten_columns: + trips_df[column] = trips[column] choices = pd.concat(choices_list) choices = choices.reindex(trips_df.index) diff --git a/activitysim/abm/models/util/logsums.py b/activitysim/abm/models/util/logsums.py index a0ec5b8457..3fe2e67d41 100644 --- a/activitysim/abm/models/util/logsums.py +++ b/activitysim/abm/models/util/logsums.py @@ -7,12 +7,12 @@ import pandas as pd from pydantic import BaseModel as PydanticBase +from activitysim.abm.models.park_and_ride_lot_choice import run_park_and_ride_lot_choice from activitysim.core import config, expressions, los, simulate, tracing, workflow from activitysim.core.configuration.logit import ( TourLocationComponentSettings, TourModeComponentSettings, ) -from activitysim.abm.models.park_and_ride_lot_choice import run_park_and_ride_lot_choice logger = logging.getLogger(__name__) @@ -179,6 +179,24 @@ def filter_chooser_columns( return choosers +def get_pnr_index_multiplier(choosers, logsum_settings): + """Keep PNR synthetic chooser IDs stable across complete chooser chunks. + + Lot choice disambiguates repeated destination-sample indices using a + multiplier rounded up to a multiple of ten. Compute it over the full + segment so changing chunk size cannot change the random-number keys. + """ + include_pnr = ( + logsum_settings.get("include_pnr_for_logsums", False) + if isinstance(logsum_settings, dict) + else getattr(logsum_settings, "include_pnr_for_logsums", False) + ) + if not include_pnr or choosers.index.is_unique: + return None + max_count = int(choosers.groupby(level=0).size().max()) + return ((max_count + 9) // 10) * 10 + + def compute_location_choice_logsums( state: workflow.State, choosers: pd.DataFrame, @@ -192,6 +210,8 @@ def compute_location_choice_logsums( in_period_col: str | None = None, out_period_col: str | None = None, duration_col: str | None = None, + explicit_chunk_size: float | None = None, + pnr_index_multiplier: int | None = None, ): """ @@ -299,6 +319,11 @@ def compute_location_choice_logsums( estimator=None, pnr_capacity_cls=None, trace_label=tracing.extend_trace_label(trace_label, "pnr_lot_choice"), + # The caller may already own the logsum chunk ledger. Propagate + # both overrides through lot choice, not just mode simulation. + chunk_size=chunk_size, + explicit_chunk_size=explicit_chunk_size, + chooser_index_multiplier=pnr_index_multiplier, ) logsum_spec = state.filesystem.read_model_spec(file_name=logsum_settings.SPEC) @@ -351,6 +376,9 @@ def compute_location_choice_logsums( trace_label=trace_label, ) + if explicit_chunk_size is None: + explicit_chunk_size = model_settings.explicit_chunk + logsums = simulate.simple_simulate_logsums( state, choosers, @@ -361,7 +389,7 @@ def compute_location_choice_logsums( chunk_size=chunk_size, chunk_tag=chunk_tag, trace_label=trace_label, - explicit_chunk_size=model_settings.explicit_chunk, + explicit_chunk_size=explicit_chunk_size, compute_settings=logsum_settings.compute_settings, ) diff --git a/activitysim/abm/models/util/tour_destination.py b/activitysim/abm/models/util/tour_destination.py index 11bdb8ecdb..068732d5f9 100644 --- a/activitysim/abm/models/util/tour_destination.py +++ b/activitysim/abm/models/util/tour_destination.py @@ -12,6 +12,7 @@ from activitysim.abm.models.util.maz_sampling import draw_maz_rands from activitysim.abm.tables.size_terms import tour_destination_size_terms from activitysim.core import ( + chunk, config, estimation, expressions, @@ -712,8 +713,7 @@ def run_destination_sample( if pre_sample_taz and not state.settings.want_dest_choice_presampling: pre_sample_taz = False logger.info( - f"Disabled destination zone presampling for {trace_label} " - f"because 'want_dest_choice_presampling' setting is False" + f"Disabled destination zone presampling for {trace_label} because 'want_dest_choice_presampling' setting is False" ) if pre_sample_taz: @@ -793,31 +793,63 @@ def run_destination_logsums( chunk_tag = "tour_destination.logsums" - # merge persons into tours - choosers = pd.merge( - destination_sample, - persons_merged, - left_on=chooser_id_column, - right_index=True, - how="left", - ) - - logger.debug("Running %s with %s rows", trace_label, len(choosers)) - state.tracing.dump_df(DUMP, persons_merged, trace_label, "persons_merged") - state.tracing.dump_df(DUMP, choosers, trace_label, "choosers") - logsums = logsum.compute_location_choice_logsums( + # One chooser row per tour, aligned with the repeated tour index of the + # sampled alternatives. Merge person attributes and run the logsum + # preprocessor inside this outer chunk so the full sampled table is never + # materialized with all person and derived columns at once. + tour_choosers = destination_sample.loc[ + ~destination_sample.index.duplicated(keep="first"), [chooser_id_column] + ] + pnr_index_multiplier = logsum.get_pnr_index_multiplier( + destination_sample, logsum_settings + ) + logsum_chunks = [] + for ( + _i, + tour_choosers_chunk, + destination_sample_chunk, + chunk_trace_label, + chunk_sizer, + ) in chunk.adaptive_chunked_choosers_and_alts( state, - choosers, - tour_purpose, - logsum_settings, - model_settings, - network_los, - chunk_size, - chunk_tag, + tour_choosers, + destination_sample, trace_label, - ) + chunk_tag, + chunk_size=chunk_size, + explicit_chunk_size=model_settings.explicit_chunk, + ): + choosers = pd.merge( + destination_sample_chunk, + persons_merged, + left_on=chooser_id_column, + right_index=True, + how="left", + ) + assert choosers.index.equals(destination_sample_chunk.index) + chunk_sizer.log_df(chunk_trace_label, "logsum_choosers", choosers) + + logsum_chunks.append( + logsum.compute_location_choice_logsums( + state, + choosers, + tour_purpose, + logsum_settings, + model_settings, + network_los, + 0, + chunk_tag, + chunk_trace_label, + explicit_chunk_size=0, + pnr_index_multiplier=pnr_index_multiplier, + ) + ) + chunk_sizer.log_df(chunk_trace_label, "logsum_choosers", None) + + logsums = pd.concat(logsum_chunks) + assert logsums.index.equals(destination_sample.index) destination_sample["mode_choice_logsum"] = logsums @@ -951,6 +983,7 @@ def run_destination_simulate( trace_choice_name="destination", estimator=estimator, skip_choice=skip_choice, + explicit_chunk_size=model_settings.explicit_chunk, compute_settings=model_settings.compute_settings, alts_context=alts_context, ) diff --git a/activitysim/abm/test/test_location_choice.py b/activitysim/abm/test/test_location_choice.py index 9bcdd7711f..c8cf1b5023 100644 --- a/activitysim/abm/test/test_location_choice.py +++ b/activitysim/abm/test/test_location_choice.py @@ -88,3 +88,63 @@ def test_estimation_override_preserves_destination_choice_logsum(monkeypatch): ) pdt.assert_frame_equal(choices, expected) assert sample is None + + +def test_location_logsums_join_person_attributes_inside_chunks(monkeypatch): + person_index = pd.Index([1, 2], name="person_id") + persons = pd.DataFrame({"income": [10, 20]}, index=person_index) + sample = pd.DataFrame( + {"alt_dest": [101, 102, 201, 202, 203]}, + index=pd.Index([1, 1, 2, 2, 2], name="person_id"), + ) + model_settings = SimpleNamespace( + LOGSUM_SETTINGS="tour_mode_choice.yaml", + LOGSUM_TOUR_PURPOSE="work", + explicit_chunk=0.5, + ) + state = SimpleNamespace(filesystem=Mock()) + chunk_sizer = Mock() + + monkeypatch.setattr( + location_choice.TourModeComponentSettings, + "read_settings_file", + lambda *args, **kwargs: SimpleNamespace(), + ) + + def chunked(*args, **kwargs): + assert args[1] is persons + assert args[2] is sample + assert kwargs["chunk_size"] == 123 + assert kwargs["explicit_chunk_size"] == 0.5 + yield 1, persons.iloc[:1], sample.iloc[:2], "logsums.i1", chunk_sizer + yield 2, persons.iloc[1:], sample.iloc[2:], "logsums.i2", chunk_sizer + + monkeypatch.setattr( + location_choice.chunk, "adaptive_chunked_choosers_and_alts", chunked + ) + chooser_lengths = [] + + def compute_logsums(_state, choosers, *args, **kwargs): + chooser_lengths.append(len(choosers)) + assert args[4] == 0 + assert kwargs["explicit_chunk_size"] == 0 + return choosers["alt_dest"] + choosers["income"] + + monkeypatch.setattr( + location_choice.logsum, "compute_location_choice_logsums", compute_logsums + ) + + result = location_choice.run_location_logsums( + state, + "work", + persons, + Mock(), + sample, + model_settings, + chunk_size=123, + chunk_tag="school_location.logsums", + trace_label="school_location.logsums.work", + ) + + assert chooser_lengths == [2, 3] + assert result[location_choice.ALT_LOGSUM].tolist() == [111, 112, 221, 222, 223] diff --git a/activitysim/abm/test/test_misc/test_tour_destination_sampling.py b/activitysim/abm/test/test_misc/test_tour_destination_sampling.py index 2157fdcadb..0dfcc0eeac 100644 --- a/activitysim/abm/test/test_misc/test_tour_destination_sampling.py +++ b/activitysim/abm/test/test_misc/test_tour_destination_sampling.py @@ -1,6 +1,7 @@ from __future__ import annotations from types import SimpleNamespace +from unittest.mock import Mock import numpy as np import pandas as pd @@ -9,6 +10,142 @@ from activitysim.core import workflow +def test_destination_logsums_join_persons_inside_chunks(monkeypatch): + persons = pd.DataFrame( + {"income": [10, 20]}, index=pd.Index([1, 2], name="person_id") + ) + sample = pd.DataFrame( + { + "person_id": [1, 1, 2, 2, 2], + "alt_dest": [101, 102, 201, 202, 203], + }, + index=pd.Index([11, 11, 22, 22, 22], name="tour_id"), + ) + model_settings = SimpleNamespace( + LOGSUM_SETTINGS="tour_mode_choice.yaml", + CHOOSER_ID_COLUMN="person_id", + explicit_chunk=0.5, + ) + state = SimpleNamespace( + filesystem=SimpleNamespace( + read_model_settings=lambda *args, **kwargs: SimpleNamespace() + ), + tracing=SimpleNamespace(dump_df=lambda *args, **kwargs: None), + ) + chunk_sizer = Mock() + + def chunked(*args, **kwargs): + assert args[1].index.tolist() == [11, 22] + assert args[2] is sample + assert kwargs["chunk_size"] == 123 + assert kwargs["explicit_chunk_size"] == 0.5 + yield 1, args[1].iloc[:1], sample.iloc[:2], "logsums.i1", chunk_sizer + yield 2, args[1].iloc[1:], sample.iloc[2:], "logsums.i2", chunk_sizer + + monkeypatch.setattr( + tour_destination.chunk, "adaptive_chunked_choosers_and_alts", chunked + ) + chooser_lengths = [] + + def compute_logsums(_state, choosers, *args, **kwargs): + chooser_lengths.append(len(choosers)) + assert args[4] == 0 + assert kwargs["explicit_chunk_size"] == 0 + return choosers["alt_dest"] + choosers["income"] + + monkeypatch.setattr( + tour_destination.logsum, + "compute_location_choice_logsums", + compute_logsums, + ) + + result = tour_destination.run_destination_logsums( + state, + "shopping", + persons, + sample, + model_settings, + Mock(), + chunk_size=123, + trace_label="non_mandatory.shopping.logsums", + ) + + assert chooser_lengths == [2, 3] + assert result["mode_choice_logsum"].tolist() == [111, 112, 221, 222, 223] + + +def test_destination_simulate_forwards_explicit_chunk_size(monkeypatch): + tours = pd.DataFrame( + {"person_id": [1, 2], "home_zone_id": [10, 20]}, + index=pd.Index([11, 22], name="tour_id"), + ) + persons = pd.DataFrame( + {"income": [100, 200]}, index=pd.Index([1, 2], name="person_id") + ) + sample = pd.DataFrame( + {"alt_dest": [101, 102, 201]}, + index=pd.Index([11, 11, 22], name="tour_id"), + ) + destination_size_terms = pd.DataFrame( + {"size_term": [1.0, 2.0, 3.0]}, + index=pd.Index([101, 102, 201], name="alt_dest"), + ) + model_settings = SimpleNamespace( + SPEC="destination.csv", + COEFFICIENTS="coefficients.csv", + CHOOSER_ID_COLUMN="person_id", + ALT_DEST_COL_NAME="alt_dest", + CHOOSER_ORIG_COL_NAME="home_zone_id", + CONSTANTS=None, + explicit_chunk=0.25, + compute_settings=SimpleNamespace(), + ) + state = SimpleNamespace( + settings=SimpleNamespace(log_alt_losers=False, use_explicit_error_terms=False), + tracing=SimpleNamespace(dump_df=lambda *args, **kwargs: None), + ) + network_los = SimpleNamespace(get_default_skim_dict=lambda: _DummySkimDict()) + monkeypatch.setattr( + tour_destination.simulate, + "spec_for_segment", + lambda *args, **kwargs: pd.DataFrame({"coefficient": [1.0]}), + ) + monkeypatch.setattr( + tour_destination.expressions, + "annotate_preprocessors", + lambda *args, **kwargs: None, + ) + captured = {} + + def simulate_sampled(*args, **kwargs): + captured.update(kwargs) + return pd.DataFrame( + {"choice": [101, 201], "logsum": [1.0, 2.0]}, index=tours.index + ) + + monkeypatch.setattr( + tour_destination, "interaction_sample_simulate", simulate_sampled + ) + + result = tour_destination.run_destination_simulate( + state, + "shopping", + tours, + persons, + sample, + want_logsums=True, + model_settings=model_settings, + network_los=network_los, + destination_size_terms=destination_size_terms, + estimator=None, + chunk_size=0, + trace_label="non_mandatory.shopping.simulate", + ) + + assert captured["explicit_chunk_size"] == 0.25 + assert result["choice"].tolist() == [101, 201] + + class _DummySkimDict: def wrap(self, orig_key, dest_key): return type("WrappedSkims", (), {"orig_key": orig_key, "dest_key": dest_key})() diff --git a/activitysim/abm/test/test_park_and_ride_logsums.py b/activitysim/abm/test/test_park_and_ride_logsums.py new file mode 100644 index 0000000000..8f9a3ac0ab --- /dev/null +++ b/activitysim/abm/test/test_park_and_ride_logsums.py @@ -0,0 +1,237 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import pandas as pd +import pytest + +from activitysim.abm.models import location_choice, park_and_ride_lot_choice as pnr +from activitysim.abm.models.util import logsums, tour_destination +from activitysim.core import interaction_simulate, workflow + + +@pytest.fixture +def pnr_context(monkeypatch): + state = workflow.State.make_temp() + state.settings.chunk_size = 1_000_000_000 + state.settings.sharrow = False + state.rng().begin_step("destination") + land_use = pd.DataFrame({"pnr_spaces": [10, 10]}, index=[1, 2]) + state.add_table("land_use", land_use) + settings = SimpleNamespace( + SPEC="unused.csv", + CONSTANTS={}, + LANDUSE_PNR_SPACES_COLUMN="pnr_spaces", + preprocessor=None, + CHOOSER_FILTER_EXPR=None, + explicit_chunk=0.5, + compute_settings=None, + ) + spec = pd.DataFrame({"coefficient": [1.0]}, index=["1"]) + fs = type(state.filesystem) + monkeypatch.setattr(fs, "read_model_spec", lambda *a, **k: spec) + monkeypatch.setattr(fs, "read_model_coefficients", lambda *a, **k: {}) + monkeypatch.setattr(fs, "get_segment_coefficients", lambda *a, **k: {}) + monkeypatch.setattr( + pnr.ParkAndRideLotChoiceSettings, "read_settings_file", lambda *a, **k: settings + ) + monkeypatch.setattr( + pnr.simulate, "eval_coefficients", lambda state, spec, *a, **k: spec + ) + monkeypatch.setattr( + pnr, + "filter_chooser_to_transit_accessible_destinations", + lambda state, choosers, *a: choosers, + ) + monkeypatch.setattr(logsums, "setup_skims", lambda *a, **k: {}) + monkeypatch.setattr(pnr.expressions, "annotate_preprocessors", lambda *a, **k: None) + monkeypatch.setattr(pnr.util, "drop_unused_columns", lambda df, *a, **k: df) + network = SimpleNamespace( + skim_time_period_label=lambda period, as_cat, broadcast_to: pd.Series( + "AM", index=broadcast_to + ) + ) + return state, land_use, network, settings + + +@pytest.mark.parametrize("mode", ["training", "adaptive", "explicit"]) +@pytest.mark.parametrize("component", ["tour", "location"]) +def test_pnr_logsums_share_outer_chunk_boundary( + pnr_context, monkeypatch, mode, component +): + state, land_use, network, pnr_settings = pnr_context + state.settings.chunk_training_mode = mode + persons = pd.DataFrame( + {"home_zone_id": [1, 2]}, index=pd.Index([1, 2], name="person_id") + ) + sample = pd.DataFrame( + {"person_id": [1, 1, 2, 2], "alt_dest": [1, 2, 1, 2]}, + index=pd.Index( + [1, 1, 2, 2], name="person_id" if component == "location" else "tour_id" + ), + ) + if component == "location": + sample = sample.drop(columns="person_id") + model = SimpleNamespace( + LOGSUM_SETTINGS="unused.yaml", + LOGSUM_TOUR_PURPOSE="work", + CHOOSER_ID_COLUMN="person_id", + explicit_chunk=0.5, + CHOOSER_ORIG_COL_NAME="home_zone_id", + ALT_DEST_COL_NAME="alt_dest", + IN_PERIOD=17, + OUT_PERIOD=8, + LOGSUM_PREPROCESSOR="preprocessor", + ) + logsum_settings = SimpleNamespace( + include_pnr_for_logsums=True, + SPEC="unused.csv", + preprocessor=None, + compute_settings=None, + ) + monkeypatch.setattr( + type(state.filesystem), "read_model_settings", lambda *a, **k: logsum_settings + ) + monkeypatch.setattr( + location_choice.TourModeComponentSettings, + "read_settings_file", + lambda *a, **k: logsum_settings, + ) + monkeypatch.setattr(logsums.config, "get_logit_model_settings", lambda *a: None) + monkeypatch.setattr(logsums.config, "get_model_constants", lambda *a: {}) + monkeypatch.setattr( + logsums.simulate, + "simple_simulate_logsums", + lambda state, choosers, *a, **k: choosers.pnr_zone_id.astype(float), + ) + sizes = [] + + def evaluate(state, choosers, *args, **kwargs): + sizes.append(len(choosers)) + return pd.Series(1, index=choosers.index) + + # Keep the actual interaction simulator and its ChunkSizer, replacing only + # utility evaluation so the test exercises the full ledger call chain. + monkeypatch.setattr(interaction_simulate, "_interaction_simulate", evaluate) + if component == "tour": + result = tour_destination.run_destination_logsums( + state, + "work", + persons, + sample, + model, + network, + state.settings.chunk_size, + "test", + ) + else: + result = location_choice.run_location_logsums( + state, + "work", + persons, + network, + sample, + model, + state.settings.chunk_size, + "logsums", + "test", + ) + assert result.mode_choice_logsum.tolist() == [1.0] * 4 + assert sum(sizes) == 4 + if mode == "explicit": + assert sizes == [2, 2] # PNR's fractional setting must not split these again. + assert pnr_settings.explicit_chunk == 0.5 + assert not state.chunk.CHUNK_SIZERS + assert not state.chunk.CHUNK_LEDGERS + + +@pytest.mark.parametrize("component", ["tour", "location"]) +def test_pnr_random_draws_and_logsums_do_not_depend_on_chunk_size( + pnr_context, monkeypatch, component +): + state, land_use, network, pnr_settings = pnr_context + state.settings.chunk_training_mode = "explicit" + persons = pd.DataFrame( + {"home_zone_id": [1, 2, 1]}, index=pd.Index([11, 22, 33], name="person_id") + ) + # Different multiplicities cross a rounding boundary; the final chooser + # also exercises a chunk whose index is unique despite the segment's not + # being unique. Every segment uses the same canonical PNR random channel. + ids = [11] * 11 + [22] * 2 + [33] + sample = pd.DataFrame( + {"person_id": ids, "alt_dest": list(range(11)) + [1, 2, 1]}, + index=pd.Index(ids, name="person_id" if component == "location" else "tour_id"), + ) + if component == "location": + sample = sample.drop(columns="person_id") + model = SimpleNamespace( + LOGSUM_SETTINGS="unused.yaml", + LOGSUM_TOUR_PURPOSE="work", + CHOOSER_ID_COLUMN="person_id", + explicit_chunk=0, + CHOOSER_ORIG_COL_NAME="home_zone_id", + ALT_DEST_COL_NAME="alt_dest", + IN_PERIOD=17, + OUT_PERIOD=8, + LOGSUM_PREPROCESSOR="preprocessor", + ) + logsum_settings = SimpleNamespace( + include_pnr_for_logsums=True, + SPEC="unused.csv", + preprocessor=None, + compute_settings=None, + ) + monkeypatch.setattr( + type(state.filesystem), "read_model_settings", lambda *a, **k: logsum_settings + ) + monkeypatch.setattr( + location_choice.TourModeComponentSettings, + "read_settings_file", + lambda *a, **k: logsum_settings, + ) + monkeypatch.setattr(logsums.config, "get_logit_model_settings", lambda *a: None) + monkeypatch.setattr(logsums.config, "get_model_constants", lambda *a: {}) + monkeypatch.setattr( + logsums.simulate, + "simple_simulate_logsums", + lambda state, choosers, *a, **k: choosers.pnr_zone_id.astype(float), + ) + draws = [] + + def evaluate(state, choosers, *args, **kwargs): + rands = state.rng().random_for_df(choosers).ravel() + draws.extend(zip(choosers.index, rands)) + return pd.Series(1 + (rands > 0.5).astype(int), index=choosers.index) + + monkeypatch.setattr(interaction_simulate, "_interaction_simulate", evaluate) + baseline = baseline_draws = None + for size in [0, 0.5, 1, 2]: + model.explicit_chunk = size + draws.clear() + if component == "tour": + result = tour_destination.run_destination_logsums( + state, "work", persons, sample.copy(), model, network, 0, "test" + ) + else: + result = location_choice.run_location_logsums( + state, + "work", + persons, + network, + sample.copy(), + model, + 0, + "logsums", + "test", + ) + if baseline is None: + baseline, baseline_draws = result, list(draws) + assert [idx for idx, rand in draws] == list(range(220, 231)) + [ + 440, + 441, + 660, + ] + else: + pd.testing.assert_frame_equal(result, baseline) + assert draws == baseline_draws + assert "pnr_lot_choice" not in state.rng().channels diff --git a/activitysim/abm/test/test_trip_mode_choice.py b/activitysim/abm/test/test_trip_mode_choice.py new file mode 100644 index 0000000000..f40a7dc7f4 --- /dev/null +++ b/activitysim/abm/test/test_trip_mode_choice.py @@ -0,0 +1,171 @@ +from __future__ import annotations + +import weakref +from contextlib import nullcontext +from types import SimpleNamespace + +import pandas as pd +import pytest + +from activitysim.abm.models import trip_mode_choice as trip_mode_choice_module + + +class _DummySkimWrapper: + df = None + + def set_df(self, df): + self.df = df + + +class _DummySkimDict: + def wrap_3d(self, **_kwargs): + return _DummySkimWrapper() + + def wrap(self, *_args): + return _DummySkimWrapper() + + def map_time_periods_from_series(self, periods): + return periods.map({"AM": 0, "PM": 1}) + + +class _DummyState: + current_model_name = "test_trip_mode_choice" + + def __init__(self, trips): + self.settings = SimpleNamespace( + downcast_int=False, + downcast_float=False, + skip_failed_choices=False, + trace_hh_id=None, + ) + self.filesystem = SimpleNamespace( + read_model_spec=lambda **_kwargs: pd.DataFrame(), + get_segment_coefficients=lambda *_args: {}, + ) + self.tables = {"trips": trips} + + def add_table(self, name, df): + self.tables[name] = df + + def get_dataframe(self, name, columns=None, as_copy=True): + df = self.tables[name] + if columns is not None: + df = df[columns] + return df.copy() if as_copy else df + + def is_table(self, _name): + return False + + +@pytest.mark.parametrize("keep_trip_period", [False, True]) +@pytest.mark.parametrize("copy_annotation_table", [False, True]) +@pytest.mark.parametrize("annotation_error", [False, True]) +def test_post_choice_annotations_preserve_requested_trip_period( + monkeypatch, keep_trip_period, copy_annotation_table, annotation_error +): + trips = pd.DataFrame( + { + "tour_id": [11, 12, 13], + "household_id": [1, 2, 3], + "primary_purpose": ["work", "shopping", "work"], + "depart": [8, 17, 9], + "origin": [1, 2, 3], + "destination": [2, 3, 1], + }, + index=pd.Index([101, 102, 103], name="trip_id"), + ) + state = _DummyState(trips) + skim_dict = _DummySkimDict() + network_los = SimpleNamespace( + skim_time_periods=SimpleNamespace(period_minutes=60), + skim_time_period_label=lambda depart, as_cat=True: depart.map( + lambda value: "AM" if value < 12 else "PM" + ), + get_default_skim_dict=lambda: skim_dict, + ) + model_settings = SimpleNamespace( + MODE_CHOICE_LOGSUM_COLUMN_NAME="mode_choice_logsum", + TOURS_MERGED_CHOOSER_COLUMNS=[], + CHOOSER_COLS_TO_KEEP=["trip_period"] if keep_trip_period else [], + FORCE_ESCORTEE_CHAUFFEUR_MODE_MATCH=False, + SPEC="trip_mode_choice.csv", + explicit_chunk=None, + compute_settings=None, + ) + + monkeypatch.setattr( + trip_mode_choice_module.tracing, "print_summary", lambda *_args, **_kwargs: None + ) + monkeypatch.setattr( + trip_mode_choice_module.config, "get_model_constants", lambda *_args: {} + ) + monkeypatch.setattr( + trip_mode_choice_module.config, "get_logit_model_settings", lambda *_args: None + ) + monkeypatch.setattr( + trip_mode_choice_module.simulate, + "eval_coefficients", + lambda _state, spec, *_args: spec, + ) + monkeypatch.setattr( + trip_mode_choice_module.simulate, + "eval_nest_coefficients", + lambda *_args: None, + ) + monkeypatch.setattr( + trip_mode_choice_module.expressions, + "annotate_preprocessors", + lambda *_args, **_kwargs: None, + ) + + chooser_refs = [] + wrappers = [] + + def choose_mode(_state, choosers, **_kwargs): + assert all(ref() is None for ref in chooser_refs) + chooser_refs.append(weakref.ref(choosers)) + wrappers[:] = _kwargs["skims"].values() + trip_mode_choice_module.simulate.set_skim_wrapper_targets( + choosers, _kwargs["skims"] + ) + return pd.DataFrame( + { + "trip_mode": "DRIVE", + "mode_choice_logsum": 1.0, + }, + index=choosers.index, + ) + + monkeypatch.setattr(trip_mode_choice_module, "mode_choice_simulate", choose_mode) + + def annotate_tables(_state, **_kwargs): + annotated = _state.get_dataframe("trips", as_copy=copy_annotation_table) + assert annotated.index.equals(trips.index) + assert annotated["trip_period"].tolist() == [0, 1, 0] + annotated["post_choice_skim_value"] = [10.0, 20.0, 30.0] + _state.add_table("trips", annotated) + if annotation_error: + raise RuntimeError("annotation failed") + + monkeypatch.setattr( + trip_mode_choice_module.expressions, "annotate_tables", annotate_tables + ) + + with pytest.raises( + RuntimeError, match="annotation failed" + ) if annotation_error else nullcontext(): + trip_mode_choice_module.trip_mode_choice( + state, + trips, + network_los, + model_settings=model_settings, + ) + + assert all(ref() is None for ref in chooser_refs) + assert all(wrapper.df.empty for wrapper in wrappers) + result = state.get_dataframe("trips", as_copy=False) + assert ("trip_period" in trips) == keep_trip_period + assert ("trip_period" in result) == keep_trip_period + if keep_trip_period: + assert result["trip_period"].tolist() == [0, 1, 0] + assert result["post_choice_skim_value"].tolist() == [10.0, 20.0, 30.0] diff --git a/activitysim/abm/test/trip_dest/test_trip_destination.py b/activitysim/abm/test/trip_dest/test_trip_destination.py index fdfd9b6de9..f3ae7ca64f 100644 --- a/activitysim/abm/test/trip_dest/test_trip_destination.py +++ b/activitysim/abm/test/trip_dest/test_trip_destination.py @@ -1,26 +1,49 @@ from __future__ import annotations import shutil +from unittest.mock import Mock from pathlib import Path import pandas as pd +import pytest + +from activitysim.abm.models import trip_destination as td +from activitysim.core import chunk from activitysim import abm # noqa: F401 from activitysim.core import workflow as wf -def test_trip_destination(tmp_path: Path): +def run_trip_destination( + tmp_path: Path, + explicit_chunk: float | None = None, + repeat_work_tours: int = 1, + chunk_training_mode: str = chunk.MODE_EXPLICIT, +): shutil.copytree( Path(__file__).parent.joinpath("configs"), tmp_path.joinpath("configs") ) shutil.copytree(Path(__file__).parent.joinpath("data"), tmp_path.joinpath("data")) + if explicit_chunk is not None: + with (tmp_path / "configs" / "trip_destination.yaml").open("a") as stream: + stream.write(f"\nexplicit_chunk: {explicit_chunk}\n") + state = wf.State.make_default(working_dir=tmp_path) + state.settings.chunk_training_mode = chunk_training_mode + if chunk_training_mode != chunk.MODE_EXPLICIT: + state.settings.chunk_size = 1_000_000 + # init tours tours = pd.read_csv( tmp_path / state.filesystem.data_dir[0] / "tours.csv" ).set_index("tour_id") + base_tour = tours.loc[[500]] + tours = pd.concat( + [tours] + + [base_tour.rename(index={500: 510 + i}) for i in range(1, repeat_work_tours)] + ).sort_index() state.add_table("tours", tours) state.tracing.register_traceable_table("tours", tours) state.get_rn_generator().add_channel("tours", tours) @@ -29,13 +52,25 @@ def test_trip_destination(tmp_path: Path): trips = pd.read_csv( tmp_path / state.filesystem.data_dir[0] / "trips.csv" ).set_index("trip_id") + base_trips = trips[trips.tour_id == 500] + repeated_trips = [trips] + for i in range(1, repeat_work_tours): + trip_copy = base_trips.copy() + trip_copy.index += 100_000 * i + trip_copy["tour_id"] = 510 + i + repeated_trips.append(trip_copy) + trips = pd.concat(repeated_trips).sort_index() state.add_table("trips", trips) state.tracing.register_traceable_table("trips", trips) state.get_rn_generator().add_channel("trips", trips) state.run.all() - out_trips = state.get_dataframe("trips") + return state.get_dataframe("trips") + + +def test_trip_destination(tmp_path: Path): + out_trips = run_trip_destination(tmp_path) # logsums are generated for intermediate trips only assert out_trips["destination_logsum"].isna().tolist() == [ @@ -46,3 +81,89 @@ def test_trip_destination(tmp_path: Path): False, True, ] + + +def test_trip_destination_chunked_logsums_match_unchunked(tmp_path: Path): + # Repeat one tour so each work-purpose presample has enough chooser rows to + # exercise the outer TAZ-to-MAZ pipeline chunker. + unchunked = run_trip_destination(tmp_path / "unchunked", repeat_work_tours=4) + chunked = run_trip_destination( + tmp_path / "chunked", + explicit_chunk=0.5, + repeat_work_tours=4, + ) + + pd.testing.assert_frame_equal(chunked, unchunked) + + +@pytest.mark.parametrize("mode", chunk.TRAINING_MODES) +@pytest.mark.parametrize("explicit_chunk", [0, 0.5, 2]) +@pytest.mark.parametrize("legacy_reason", [None, "estimator", "sample_table"]) +def test_destination_pipeline_chunk_boundary( + tmp_path, monkeypatch, mode, explicit_chunk, legacy_reason +): + state = wf.State.make_default( + working_dir=tmp_path, + configs_dir=Path(__file__).parent / "configs", + data_dir=Path(__file__).parent / "data", + ) + state.settings.chunk_training_mode = mode + state.settings.chunk_size = 1_000_000 + settings = td.TripDestinationSettings.model_construct(explicit_chunk=explicit_chunk) + trips = pd.DataFrame({"value": range(4)}, index=pd.Index(range(4), name="trip_id")) + calls = [] + + def choose( + _state, purpose, chooser, alternatives, tours, model_settings, *args, **kwargs + ): + calls.append((chooser.index.tolist(), model_settings)) + return chooser.copy(), None + + monkeypatch.setattr(td, "_choose_trip_destination_unchunked", choose) + monkeypatch.setattr(td.mem, "release_memory", lambda: False) + use_outer = mode == chunk.MODE_EXPLICIT and explicit_chunk and legacy_reason is None + if not use_outer: + outer = Mock( + side_effect=AssertionError("legacy path must not create an outer chunker") + ) + monkeypatch.setattr(td.chunk, "adaptive_chunked_choosers", outer) + result, sample = td.choose_trip_destination( + state, + "work", + trips, + None, + None, + settings, + False, + legacy_reason == "sample_table", + None, + None, + object() if legacy_reason == "estimator" else None, + 1_000_000, + "test", + ) + pd.testing.assert_frame_equal(result, trips) + assert sample is None + assert settings.explicit_chunk == explicit_chunk + if use_outer: + assert [rows for rows, _ in calls] == [[0, 1], [2, 3]] + assert all( + inner.explicit_chunk == 0 and inner is not settings for _, inner in calls + ) + else: + assert calls == [([0, 1, 2, 3], settings)] + + +def test_trip_destination_training_ignores_explicit_chunk(tmp_path): + unchunked = run_trip_destination( + tmp_path / "default", + repeat_work_tours=4, + chunk_training_mode=chunk.MODE_RETRAIN, + ) + explicit = run_trip_destination( + tmp_path / "explicit", + explicit_chunk=0.5, + repeat_work_tours=4, + chunk_training_mode=chunk.MODE_RETRAIN, + ) + pd.testing.assert_frame_equal(explicit, unchunked) diff --git a/activitysim/core/interaction_simulate.py b/activitysim/core/interaction_simulate.py index 496abf79e0..c36791179c 100644 --- a/activitysim/core/interaction_simulate.py +++ b/activitysim/core/interaction_simulate.py @@ -993,6 +993,7 @@ def interaction_simulate( estimator=None, explicit_chunk_size=0, compute_settings: ComputeSettings | None = None, + chunk_size: int | None = None, ): """ Run a simulation in the situation in which alternatives must @@ -1032,6 +1033,9 @@ def interaction_simulate( when household tracing enabled. No tracing occurs if label is empty or None. trace_choice_name: str This is the column label to be used in trace file csv dump of choices + chunk_size : int, optional + Adaptive memory budget; zero disables inner chunking. Defaults to the + global setting when omitted. explicit_chunk_size : float, optional If > 0, specifies the chunk size to use when chunking the interaction simulation. If < 1, specifies the fraction of the total number of choosers. @@ -1054,7 +1058,11 @@ def interaction_simulate( chunk_trace_label, chunk_sizer, ) in chunk.adaptive_chunked_choosers( - state, choosers, trace_label, explicit_chunk_size=explicit_chunk_size + state, + choosers, + trace_label, + explicit_chunk_size=explicit_chunk_size, + chunk_size=chunk_size, ): choices = _interaction_simulate( state, diff --git a/activitysim/core/mem.py b/activitysim/core/mem.py index fbfa11fbe5..19440f24cd 100644 --- a/activitysim/core/mem.py +++ b/activitysim/core/mem.py @@ -2,12 +2,14 @@ # See full license in LICENSE.txt. from __future__ import annotations +import ctypes import datetime import gc import glob import logging import multiprocessing import os +import sys import threading import time @@ -15,7 +17,7 @@ import pandas as pd import psutil -from activitysim.core import config, util, workflow +from activitysim.core import util, workflow logger = logging.getLogger(__name__) @@ -29,13 +31,79 @@ MEM_TICK = 0 MEM_LOG_FILE_NAME = "mem.csv" -OMNIBUS_LOG_FILE_NAME = f"omnibus_mem.csv" +OMNIBUS_LOG_FILE_NAME = "omnibus_mem.csv" SUMMARY_BIN_SIZE_IN_SECONDS = 15 mem_log_lock = threading.Lock() +def release_memory(): + """Return unused allocator pages to the operating system when possible. + + Garbage collection destroys unreachable Python objects, but native arrays + created by pandas, NumPy, and compiled model evaluators can leave free pages + in the process allocator. Long, segmented model steps may therefore retain + a high resident set even after each segment's temporary frames are gone. + + The allocator-pressure calls below are advisory and platform specific. A + failed or unavailable call is harmless: collection has still occurred and + the model continues normally. + """ + + was_disabled = not gc.isenabled() + if was_disabled: + gc.enable() + gc.collect() + if was_disabled: + gc.disable() + + try: + if sys.platform.startswith("linux"): + libc = ctypes.CDLL(None) + malloc_trim = getattr(libc, "malloc_trim", None) + if malloc_trim is not None: + malloc_trim.argtypes = [ctypes.c_size_t] + malloc_trim.restype = ctypes.c_int + return bool(malloc_trim(0)) + + elif sys.platform == "darwin": + libc = ctypes.CDLL(None) + malloc_default_zone = getattr(libc, "malloc_default_zone", None) + pressure_relief = getattr(libc, "malloc_zone_pressure_relief", None) + if malloc_default_zone is not None and pressure_relief is not None: + malloc_default_zone.argtypes = [] + malloc_default_zone.restype = ctypes.c_void_p + pressure_relief.argtypes = [ctypes.c_void_p, ctypes.c_size_t] + pressure_relief.restype = ctypes.c_size_t + return bool(pressure_relief(malloc_default_zone(), 0)) + + elif sys.platform == "win32": + kernel32 = ctypes.WinDLL("kernel32", use_last_error=True) + get_current_process = kernel32.GetCurrentProcess + get_current_process.argtypes = [] + get_current_process.restype = ctypes.c_void_p + set_working_set_size = kernel32.SetProcessWorkingSetSize + set_working_set_size.argtypes = [ + ctypes.c_void_p, + ctypes.c_size_t, + ctypes.c_size_t, + ] + set_working_set_size.restype = ctypes.c_int + maximum_size = ctypes.c_size_t(-1).value + return bool( + set_working_set_size( + get_current_process(), + maximum_size, + maximum_size, + ) + ) + except (AttributeError, OSError, TypeError, ValueError): + logger.debug("Platform allocator did not release memory", exc_info=True) + + return False + + def time_bin(timestamps): bins_size_in_seconds = SUMMARY_BIN_SIZE_IN_SECONDS epoch = pd.Timestamp("1970-01-01") @@ -217,7 +285,7 @@ def trace_memory_info(event, trace_ticks=0, force_garbage_collect=False, *, stat child_info = child.memory_info() full_rss += child_info.rss num_children += 1 - except (psutil.NoSuchProcess, psutil.AccessDenied) as e: + except (psutil.NoSuchProcess, psutil.AccessDenied): pass noteworthy = ( @@ -296,7 +364,7 @@ def shared_memory_size(data_buffers): if data_buffers is None: data_buffers = {} - for k, data_buffer in data_buffers.items(): + for _k, data_buffer in data_buffers.items(): if isinstance(data_buffer, str) and data_buffer.startswith("sh.Dataset:"): from sharrow import Dataset diff --git a/activitysim/core/test/test_mem.py b/activitysim/core/test/test_mem.py new file mode 100644 index 0000000000..f1e87f6390 --- /dev/null +++ b/activitysim/core/test/test_mem.py @@ -0,0 +1,87 @@ +from __future__ import annotations + +import ctypes +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from activitysim.core import mem + + +def test_release_memory_is_advisory(): + assert isinstance(mem.release_memory(), bool) + + +@pytest.mark.parametrize("platform", ["linux", "darwin", "win32", "unsupported"]) +@pytest.mark.parametrize("gc_enabled", [True, False]) +@pytest.mark.parametrize( + "outcome", ["success", "no_release", "missing", "load_error", "call_error"] +) +def test_release_memory_platforms(monkeypatch, platform, gc_enabled, outcome): + # Mock the platform module itself so pytest and dependencies retain the host + # platform, and never invoke a foreign allocator on the test machine. + monkeypatch.setattr(mem, "sys", SimpleNamespace(platform=platform)) + gc = Mock() + gc.isenabled.return_value = gc_enabled + monkeypatch.setattr(mem, "gc", gc) + release = Mock(return_value=1 if outcome == "success" else 0) + handle = Mock(return_value=1234) + library = SimpleNamespace() + if outcome != "missing": + library = SimpleNamespace( + malloc_trim=release, + malloc_default_zone=handle, + malloc_zone_pressure_relief=release, + GetCurrentProcess=handle, + SetProcessWorkingSetSize=release, + ) + loader = Mock(return_value=library) + if outcome == "load_error": + loader.side_effect = OSError("allocator unavailable") + elif outcome == "call_error": + release.side_effect = OSError("allocator call failed") + monkeypatch.setattr(mem.ctypes, "CDLL", loader) + monkeypatch.setattr(mem.ctypes, "WinDLL", loader, raising=False) + + assert mem.release_memory() is (outcome == "success" and platform != "unsupported") + gc.collect.assert_called_once_with() + if gc_enabled: + gc.enable.assert_not_called() + gc.disable.assert_not_called() + else: + assert [call[0] for call in gc.mock_calls] == [ + "isenabled", + "enable", + "collect", + "disable", + ] + if platform == "unsupported": + loader.assert_not_called() + else: + if platform == "win32": + loader.assert_called_once_with("kernel32", use_last_error=True) + else: + loader.assert_called_once_with(None) + if outcome in ("success", "no_release", "call_error"): + if platform == "linux": + release.assert_called_once_with(0) + assert release.argtypes == [ctypes.c_size_t] + assert release.restype == ctypes.c_int + elif platform == "darwin": + handle.assert_called_once_with() + release.assert_called_once_with(1234, 0) + assert handle.restype == ctypes.c_void_p + assert release.argtypes == [ctypes.c_void_p, ctypes.c_size_t] + assert release.restype == ctypes.c_size_t + else: + handle.assert_called_once_with() + maximum_size = ctypes.c_size_t(-1).value + release.assert_called_once_with(1234, maximum_size, maximum_size) + assert handle.restype == ctypes.c_void_p + assert release.argtypes == [ + ctypes.c_void_p, + ctypes.c_size_t, + ctypes.c_size_t, + ] + assert release.restype == ctypes.c_int diff --git a/activitysim/core/test/test_workflow_runner.py b/activitysim/core/test/test_workflow_runner.py new file mode 100644 index 0000000000..89c01ec3a0 --- /dev/null +++ b/activitysim/core/test/test_workflow_runner.py @@ -0,0 +1,37 @@ +from __future__ import annotations + +from types import SimpleNamespace + +from activitysim.core import mp_tasks +from activitysim.core.run_id import RunId +from activitysim.core.workflow.runner import Runner + + +class ProgrammaticState: + def __init__(self): + self.settings = SimpleNamespace(memory_profile=False, multiprocess=True) + self.tracing = SimpleNamespace(run_id=RunId("abc123")) + + def __contains__(self, key): + return key == "preload_injectables" + + def get_injectable(self, key): + assert key != "run_id" + return f"value-for-{key}" + + +def test_programmatic_multiprocess_run_gets_tracing_run_id(monkeypatch): + state = ProgrammaticState() + runner = Runner(state) + received = {} + + def capture_run_multiprocess(passed_state, injectables): + assert passed_state is state + received.update(injectables) + + monkeypatch.setattr(mp_tasks, "run_multiprocess", capture_run_multiprocess) + + runner.all(config_logger=False, filter_warnings=False) + + assert received["run_id"] == "abc123" + assert received["settings"] is state.settings diff --git a/activitysim/core/workflow/runner.py b/activitysim/core/workflow/runner.py index 79ecd0ed4f..df56f4cbc4 100644 --- a/activitysim/core/workflow/runner.py +++ b/activitysim/core/workflow/runner.py @@ -5,6 +5,7 @@ import time from collections.abc import Callable, Iterable from datetime import timedelta +from typing import Any from activitysim.core import tracing from activitysim.core.exceptions import DuplicateWorkflowNameError @@ -266,8 +267,7 @@ def _pre_run_step(self, model_name: str) -> bool | None: if self._obj.settings.duplicate_step_execution == "error": checkpointed_model_bullets = "\n - ".join(checkpointed_models) raise DuplicateWorkflowNameError( - f"Checkpointed Models:\n - {checkpointed_model_bullets}\n" - f"Cannot run model '{model_name}' more than once" + f"Checkpointed Models:\n - {checkpointed_model_bullets}\nCannot run model '{model_name}' more than once" ) self._obj.rng().begin_step(model_name) @@ -412,7 +412,19 @@ def all( from activitysim.cli.run import INJECTABLES from activitysim.core import mp_tasks - injectables = {k: self._obj.get_injectable(k) for k in INJECTABLES} + # ``run_id`` is initialized lazily by the tracing accessor and + # is not necessarily present as a top-level injectable when a + # caller constructs State programmatically. The CLI mirrors it + # explicitly in ``handle_standard_args``; do the equivalent + # here so ``State.run.all()`` works for multiprocessing too. + injectables = { + key: ( + self._obj.tracing.run_id + if key == "run_id" + else self._obj.get_injectable(key) + ) + for key in INJECTABLES + } injectables["settings"] = self._obj.settings # injectables["settings_package"] = state.settings.dict() mp_tasks.run_multiprocess(self._obj, injectables) diff --git a/pyproject.toml b/pyproject.toml index df86451c8f..4ac291fb62 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,6 +48,12 @@ exclude = [ "activitysim/abm/models/util/test*", ] +# Explicitly include the example resources used by ``activitysim create``. +# Relying on setuptools-scm's implicit file list works in a Git checkout, but +# produces an incomplete wheel in Docker and other source-only build contexts. +[tool.setuptools.package-data] +activitysim = ["examples/**/*"] + [tool.setuptools_scm] fallback_version = "999" write_to = "activitysim/_generated_version.py"