Source code for urbansim_templates.modelmanager

from __future__ import print_function

import os
import copy
import pickle
from collections import OrderedDict

import orca
from urbansim.utils import yamlio

from .__init__ import __version__
from .utils import update_name, version_greater_or_equal


_templates = {}  # global registry of template classes
_steps = {}  # global registry of model steps in memory
_disk_store = None  # path to saved steps on disk


[docs]def template(cls): """ This is a decorator for ModelManager-compliant template classes. Place `@modelmanager.template` on the line before a class defintion. This makes the class available to ModelManager (e.g. for reading saved steps from disk) whenever it's imported. """ _templates[cls.__name__] = cls return cls
[docs]def initialize(path='configs'): """ Load saved model steps from disk. Each file in the directory will be checked for compliance with the ModelManager YAML format and then loaded into memory. If run multiple times, steps will be cleared from memory and re-loaded. Parameters ---------- path : str Path to config directory, either absolute or relative to the Python working directory """ if not os.path.exists(path): print("Path not found: {}".format(os.path.join(os.getcwd(), path))) # TO DO - automatically create directory if run again after warning? return global _steps, _disk_store _steps = {} # clear memory _disk_store = path # save initialization path files = [] for f in os.listdir(path): if f[-5:] == '.yaml': files.append(os.path.join(path, f)) if len(files) == 0: print("No yaml files found in path '{}'".format(path)) return steps = [] for f in files: d = yamlio.yaml_to_dict(str_or_buffer=f) if 'modelmanager_version' in d: # TO DO - check that file name matches object name in the file? if version_greater_or_equal(d['modelmanager_version'], '0.1.dev8'): # This is the version that switched from a single file to multiple files # with one object stored in each steps.append(d) if len(steps) == 0: print("No files from ModelManager 0.1.dev8 or later found in path '{}'"\ .format(path)) for d in steps: # TO DO - check for this key, to be safe step = build_step(d['saved_object']) register(step, save_to_disk=False)
[docs]def build_step(d): """ Build a model step object from a saved dictionary. This includes loading supplemental objects from disk. Parameters ---------- d : dict Representation of a model step. Returns ------- object """ template = d['meta']['template'] if 'meta' in d else d['template'] if 'supplemental_objects' in d: for i, item in enumerate(d['supplemental_objects']): content = load_supplemental_object(d['name'], **item) d['supplemental_objects'][i]['content'] = content return _templates[template].from_dict(d)
[docs]def load_supplemental_object(step_name, name, content_type, required=True): """ Load a supplemental object from disk. Parameters ---------- step_name : str Name of the associated model step. name : str Name of the supplemental object. content_type : str Currently supports 'pickle'. required : bool, optional Whether the supplemental object is required (not yet supported). Returns ------- object """ if (content_type == 'pickle'): with open(os.path.join(_disk_store, step_name+'-'+name+'.pkl'), 'rb') as f: return pickle.load(f)
[docs]def register(step, save_to_disk=True): """ Register a model step with ModelManager and Orca. This includes saving it to disk, optionally, so it can be automatically loaded in the future. Registering a step will overwrite any previously loaded step with the same name. If a name has not yet been assigned, one will be generated from the template name and a timestamp. If the model step includes an attribute 'autorun' that's set to True, the step will run after being registered. Parameters ---------- step : object Returns ------- None """ # Currently supporting both step.name and step.meta.name if hasattr(step, 'meta'): # TO DO: move the name updating to CoreTemplateSettings? step.meta.name = update_name(step.meta.template, step.meta.name) name = step.meta.name else: step.name = update_name(step.template, step.name) name = step.name if save_to_disk: save_step_to_disk(step) print("Registering model step '{}'".format(name)) _steps[name] = step # Create a callable that runs the model step, and register it with orca def run_step(): return step.run() orca.add_step(name, run_step) if hasattr(step, 'meta'): if step.meta.autorun: orca.run([name]) elif hasattr(step, 'autorun'): if step.autorun: orca.run([name])
[docs]def list_steps(): """ Return a list of registered steps, with name, template, and tags for each. Returns ------- list of dicts, ordered by name """ steps = [] for k in sorted(_steps.keys()): if hasattr(_steps[k], 'meta'): steps += [{'name': _steps[k].meta.name, 'template': _steps[k].meta.template, 'tags': _steps[k].meta.tags, 'notes': _steps[k].meta.notes}] else: steps += [{'name': _steps[k].name, 'template': _steps[k].template, 'tags': _steps[k].tags}] return steps
[docs]def save_step_to_disk(step): """ Save a model step to disk, over-writing the previous file. The file will be named 'model-name.yaml' and will be saved to the initialization directory. """ name = step.meta.name if hasattr(step, 'meta') else step.name if _disk_store is None: print("Please run 'modelmanager.initialize()' before registering new model steps") return print("Saving '{}.yaml': {}".format(name, os.path.join(os.getcwd(), _disk_store))) d = step.to_dict() # Save supplemental objects if 'supplemental_objects' in d: for item in filter(None, d['supplemental_objects']): save_supplemental_object(name, **item) del item['content'] # Save main yaml file headers = {'modelmanager_version': __version__} content = OrderedDict(headers) content.update({'saved_object': d}) yamlio.convert_to_yaml(content, os.path.join(_disk_store, name+'.yaml'))
[docs]def save_supplemental_object(step_name, name, content, content_type, required=True): """ Save a supplemental object to disk. Parameters ---------- step_name : str Name of the associated model step. name : str Name of the supplemental object. content : obj Object to save. content_type : str Currently supports 'pickle'. required : bool, optional Whether the supplemental object is required (not yet supported). """ if content_type == 'pickle': content.to_pickle(os.path.join(_disk_store, step_name+'-'+name+'.pkl'))
[docs]def get_step(name): """ Return the class representation of a registered step, by name. Parameters ---------- name : str Returns ------- instance of a template class """ return copy.deepcopy(_steps[name])
[docs]def remove_step(name): """ Remove a model step, by name. It will immediately be removed from ModelManager and from disk, but will remain registered in Orca until the current Python process terminates. Parameters ---------- name : str """ print("Removing '{}' and '{}.yaml'".format(name, name)) d = _steps[name].to_dict() if 'supplemental_objects' in d: for item in filter(None, d['supplemental_objects']): remove_supplemental_object(name, item['name'], item['content_type']) del _steps[name] os.remove(os.path.join(_disk_store, name+'.yaml'))
[docs]def remove_supplemental_object(step_name, name, content_type): """ Remove a supplemental object from disk. Parameters ---------- step_name : str Name of the associated model step. name : str Name of the supplemental object. content_type : str Currently supports 'pickle'. """ # TO DO - check that the file exists first if content_type == 'pickle': os.remove(os.path.join(_disk_store, step_name+'-'+name+'.pkl'))
[docs]def get_config_dir(): """ Return the config directory, for other services that need to interoperate. Returns ------- str """ return _disk_store