Merge pull request #31 from SchmidtDSE/feat/better_logging_and_artifacts

Logging configurability and simulation artifacts

.devcontainer/DEVELOPING.MD

@@ -0,0 +1,73 @@
+## Logging Configuration
+
+Logging is configurable within the simulation using the following workflow.
+
+
+#### Logging Pattern
+
+First, find a reasonable location to trigger a log event - this should be associated with a particular model 'event', whether it occurs to the `agent` or to the `sim` itself. For eg, the `sim` will 'experience' a step, and during that step, each `agent` will make a step - but we likely want to log different things at the time of an agent step and a sim step.
+
+Accordingly, at the top of `vegetation/config/logging.py`, we define a events as:
+
+```
+
+class AgentEventType(Enum):
+    ON_CREATE = "on_create"
+    ON_STEP = "on_step"
+    ON_SURVIVE = "on_survive"
+    ON_DEATH = "on_death"
+    ON_TRANSITION = "on_transition"
+    ON_DISPERSE = "on_disperse"
+
+
+class SimEventType(Enum):
+    ON_START = "on_start"
+    ON_STEP = "on_step"
+    ON_END = "on_end"
+    ON_MANAGE = "on_manage"
+
+```
+
+We will use these later on to determine the format, flavor, and logging level of the log message we issue when the given event occurs. 
+
+#### JSON Configuration
+
+For each particular logging event defined, we have a representative key / value entry within `vegetation/config/logging_config.json`. For eg, for the event `AgentEventType.ON_SURVIVE`, we have defined:
+
+```json
+{
+    "agent": {
+        "JoshuaTreeAgent": {
+            "on_survive": "{STD_INDENT}💪 Agent {agent.unique_id} survived! ({agent.life_stage} w/ survival rate {survival_rate:.2f})",
+        ...
+    },
+    "sim": ...
+}
+```
+
+Note that in this example we added `...` to signify the rest of the json config for the other event types, for clarity. 
+
+There are a couple of conventions we assume (logic defined by `vegetation/config/logging.py`). These log strings will be passed to `string.format` with some custom logic - any variables defined within curly braces will be obtained in one of two ways:
+
+- First, if they are a single variable name with no `.` within the variable (for eg `survival_rate`), the logger will expect to be passed this variable within a keyword argument defined in `context` dict (more on that below).
+
+- If the variable contains a `.` - for eg `agent.life_stage`, this will be assumed to be obtained from within the agent method where the logging is triggered. The logging event has access to the agent object by default, so these don't need to be passed in `context` - but they do need to exist as a valid attribute within the `agent` class (for example, `agent.life_stage` must be valid and existing).
+
+    - _For now, this only works one level deep and doesn't support execution - for eg, `agent.indices.min_index` or `agent.indices[0]` or `int(agent.life_stage)` or similar will not work. We could investigate this if it seems useful_, though.
+
+- We also arbitrarily define some useful constants: `STD_INDENT`, `WIDE_INDENT`, `VERY_WIDE_INDENT` - these are also defined in `vegetation/config/logging.py`. This is nice because we can change the indent size globally across all logs in one spot, for eg, but might also be useful for other 'constants' that don't rely on agent state or application logic really. These can be readily ignored in favor of hard-coding spaces or tabs for example, but this logic is set up in case it proves useful. 
+
+#### Triggering a Log
+
+The agent(s) and simulation both have their own loggers - `agent_logger` and `sim_logger`, respectively. These can be accessed at any point to issue a log event according to the event type passed along, as follows:
+
+```
+self.agent_logger.log_agent_event(
+    self,
+    AgentEventType.ON_SURVIVE,
+    context={"survival_rate": survival_rate},
+)
+```
+
+Using this method call, the `agent_logger` will grab the appropriate log method defined in `logging_config.json`, according to the log event type (here `AgentEventType.ON_SURVIVE`), and pass along the agent itself (here, `self`, since this method is triggered from within the agent class method `on_step`) and any additional context necessary (here `survival_rate`, since this is ephemeral and not stored by the agent - but do we want to log it, so we pass it for just this log event).
+

.devcontainer/devcontainer.json

@@ -21,5 +21,12 @@
     "github.vscode-github-actions",
     "ms-azuretools.vscode-docker"
   ],
-  "postCreateCommand": ".devcontainer/scripts/run_post_create.sh"
+  "postCreateCommand": ".devcontainer/scripts/run_post_create.sh",
+  "customizations": {
+    "vscode": {
+      "settings": {
+        "python.defaultInterpreterPath": "${containerWorkspaceFolder}/.pixi/envs/default/bin/python"
+      }
+    }
+  }
 }

vegetation/app.py

@@ -73,7 +73,9 @@ def cell_portrayal(agent):
 # manually trigger `_on_start` for now to ensure init is the same as before,
 # but we ideally want this to not kick when the solara viz is created. I suppose
 # maybe we need an `_on_viz` hook or something?
-vegetation_model = Vegetation(bounds=TST_JOTR_BOUNDS)
+vegetation_model = Vegetation(
+    bounds=TST_JOTR_BOUNDS, log_config_path="vegetation/config/logging_config.json"
+)
 cache_manager = CacheManager(bounds=TST_JOTR_BOUNDS, epsg=4326, model=vegetation_model)
 
 cache_manager.populate_elevation_cache_if_not_exists()

vegetation/batch.py

@@ -7,9 +7,6 @@
     DEM_STAC_PATH,
 )
 
-# TODO: Batch run can't see local cache
-# Issue URL: https://github.com/SchmidtDSE/mesa_abm_poc/issues/16
-
 # TODO: Implement early stopping when all the JOTR die off
 # Issue URL: https://github.com/SchmidtDSE/mesa_abm_poc/issues/18
 

vegetation/cache_manager.py

@@ -7,6 +7,7 @@
 from functools import cached_property
 from pystac_client import Client as PystacClient
 import planetary_computer
+import logging
 
 from vegetation.config.paths import DEM_STAC_PATH, LOCAL_STAC_CACHE_FSTRING
 from vegetation.config.aoi import TST_JOTR_BOUNDS
@@ -50,14 +51,15 @@ def _docker_host_cache_paths(self) -> dict:
         return docker_host_cache_dict
 
     def get_elevation_from_stac(self):
-        print("Collecting STAC Items for elevation")
+
+        logging.debug("Collecting STAC Items for elevation")
         items_generator = self.pystac_client.search(
             collections=["cop-dem-glo-30"],
             bbox=self.bounds,
         ).items()
 
         items = [item for item in items_generator]
-        print(f"Found {len(items)} items")
+        logging.debug(f"Found {len(items)} items")
 
         elevation = stackstac.stack(
             items=items,
@@ -81,10 +83,10 @@ def populate_elevation_cache_if_not_exists(self):
         elevation_cache_path = self._cache_paths["elevation"]
 
         if os.path.exists(elevation_cache_path):
-            print(f"Local elevation cache found: {elevation_cache_path}")
+            logging.debug(f"Local elevation cache found: {elevation_cache_path}")
             return
 
-        print("No local cache found, downloading elevation from STAC")
+        logging.warning("No local cache found, downloading elevation from STAC")
         time_at_start = time.time()
 
         elevation = self.get_elevation_from_stac()
@@ -104,7 +106,7 @@ def populate_elevation_cache_if_not_exists(self):
             attr_name="elevation",
         )
 
-        print(f"Saving elevation to local cache: {elevation_cache_path}")
+        logging.debug(f"Saving elevation to local cache: {elevation_cache_path}")
         os.makedirs(os.path.dirname(elevation_cache_path), exist_ok=True)
         elevation_layer.to_file(elevation_cache_path)
 
@@ -117,12 +119,12 @@ def populate_elevation_cache_if_not_exists(self):
                 os.path.dirname(docker_host_elevation_cache_path), exist_ok=True
             )
 
-            print(
+            logging.debug(
                 f"Also saving elevation to Docker host cache (to speed up cache build later on this machine): {docker_host_elevation_cache_path}"
             )
             elevation_layer.to_file(docker_host_elevation_cache_path)
 
-        print(f"Downloaded elevation in {time.time() - time_at_start} seconds")
+        logging.debug(f"Downloaded elevation in {time.time() - time_at_start} seconds")
 
 
 if __name__ == "__main__":

vegetation/config/logging.py

@@ -1,10 +1,183 @@
+# logging.py
+
 import logging
+import json
+import string
+from enum import Enum, auto
+from typing import Dict, Optional
+
+
+class AgentEventType(Enum):
+    ON_CREATE = "on_create"
+    ON_STEP = "on_step"
+    ON_SURVIVE = "on_survive"
+    ON_DEATH = "on_death"
+    ON_TRANSITION = "on_transition"
+    ON_DISPERSE = "on_disperse"
+
+
+class SimEventType(Enum):
+    ON_START = "on_start"
+    ON_STEP = "on_step"
+    ON_END = "on_end"
+    ON_MANAGE = "on_manage"
 
-{
-    "vegetation.model.Vegetation": {
-        "__init__": {"log_level": logging.DEBUG, "indent": 0},
-        "on_step": {"log_level": logging.DEBUG, "indent": 1},
-        "add_agents_from_managment_draw": {"log_level": logging.DEBUG, "indent": 3},
-        "complete": {"log_level": logging.DEBUG, "indent": 0},
-    }
+
+STD_FORMATTERS = {
+    "STD_INDENT": "    ",
+    "WIDE_INDENT": "        ",
+    "VERY_WIDE_INDENT": "            ",
 }
+
+
+class FallbackFormatter(string.Formatter):
+    def get_field(self, field_name, args, kwargs):
+
+        try:
+            # Check if field_name contains dots (e.g., 'agent.unique_id')
+            if "." in field_name:
+                obj_name, attr = field_name.split(".", 1)
+                assert obj_name in ["agent", "sim"], f"Invalid object name: {obj_name}"
+                if hasattr(kwargs.get(obj_name), attr):
+                    attr_value = getattr(kwargs[obj_name], attr)
+                    return attr_value, field_name
+                else:
+                    raise AttributeError(
+                        f"Could not find {attr} in {obj_name}'s attributes"
+                    )
+            else:
+                return super().get_field(field_name, args, kwargs)
+        except (KeyError, AttributeError) as e:
+            raise ValueError(
+                f"Could not find {field_name} in context or agent's attributes"
+            ) from e
+
+
+class LogConfig:
+    _instance = None
+
+    @classmethod
+    def initialize(cls, config_path):
+        instance = cls()
+        instance.load_config(config_path)
+        return instance
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+            cls._instance._agent_templates = {}
+            cls._instance._sim_templates = {}
+        return cls._instance
+
+    def load_config(self, config_path: str):
+        with open(config_path, "r") as f:
+            log_config_dict = json.load(f)
+            if "agent" in log_config_dict:
+                self._agent_templates = log_config_dict["agent"]
+            if "sim" in log_config_dict:
+                self._sim_templates = log_config_dict["sim"]
+
+    def update_agent_template(self, agent_type: str, event_type: str, template: str):
+        if agent_type not in self._templates:
+            self._agent_templates[agent_type] = {}
+        self._agent_templates[agent_type][event_type] = template
+
+    def update_sim_template(self, sim_type: str, event_type: str, template: str):
+        if sim_type not in self._templates:
+            self._sim_templates[sim_type] = {}
+        self._sim_templates[sim_type][event_type] = template
+
+    def get_agent_template(self, agent_type: str, event_type: str) -> Optional[str]:
+        return self._agent_templates.get(agent_type, {}).get(event_type)
+
+    def get_sim_template(self, sim_type: str, event_type: str) -> Optional[str]:
+        return self._sim_templates.get(sim_type, {}).get(event_type)
+
+
+# TODO: Figure out if AgentLogger and SimLogger need to be different classes
+# Issue URL: https://github.com/SchmidtDSE/mesa_abm_poc/issues/32
+# It might make sense to do these in one class
+
+
+class AgentLogger:
+    _instance = None
+    _fallback_formatter = FallbackFormatter()
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+            cls._instance.config = LogConfig()
+            cls._instance._setup_logger()
+        return cls._instance
+
+    def _setup_logger(self):
+        self.logger = logging.getLogger("agent_logger")
+        self.logger.setLevel(logging.INFO)
+        self.logger.propagate = False
+
+        if not self.logger.handlers:
+            ch = logging.StreamHandler()
+            formatter = logging.Formatter("%(message)s")
+            ch.setFormatter(formatter)
+            self.logger.addHandler(ch)
+
+    def log_agent_event(self, agent, event_type: AgentEventType, context: Dict = None):
+
+        if not agent.log_level:
+            return
+
+        if context is None:
+            context = {}
+
+        template = self.config.get_agent_template(
+            agent.__class__.__name__, event_type.value
+        )
+
+        if template:
+            context["agent"] = agent
+            context.update(STD_FORMATTERS)
+            message = self._fallback_formatter.format(template, **context)
+            self.logger.log(agent.log_level, message)
+
+
+class SimLogger:
+    _instance = None
+    _fallback_formatter = FallbackFormatter()
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+            cls._instance.config = LogConfig()
+            cls._instance._setup_logger()
+        return cls._instance
+
+    def _setup_logger(self):
+        self.logger = logging.getLogger("sim_logger")
+        self.logger.setLevel(logging.INFO)
+        self.logger.propagate = False
+
+        if not self.logger.handlers:
+            ch = logging.StreamHandler()
+            formatter = logging.Formatter("%(message)s")
+            ch.setFormatter(formatter)
+            self.logger.addHandler(ch)
+
+    def log_sim_event(
+        self, sim, event_type: SimEventType, context: Dict = None, level=logging.INFO
+    ):
+
+        if not sim.log_level:
+            return
+
+        if context is None:
+            context = {}
+
+        template = self.config.get_sim_template(
+            sim.__class__.__name__, event_type.value
+        )
+
+        if template:
+            context["sim"] = sim
+            context.update(STD_FORMATTERS)
+            message = self._fallback_formatter.format(template, **context)
+            self.logger.log(sim.log_level, message)

vegetation/config/logging_config.json

@@ -0,0 +1,19 @@
+
+{
+    "agent": {
+        "JoshuaTreeAgent": {
+            "on_create": "{VERY_WIDE_INDENT}🌱 Agent {agent.unique_id} created at {agent.indices}",
+            "on_survive": "{STD_INDENT}💪 Agent {agent.unique_id} survived! ({agent.life_stage} w/ survival rate {survival_rate:.2f})",
+            "on_death": "{STD_INDENT}💀 Agent {agent.unique_id} died! ({agent.life_stage} w/ survival rate {survival_rate:.2f})",
+            "on_transition": "{WIDE_INDENT}🔄 Agent {agent.unique_id} promoted to {agent.life_stage}",
+            "on_disperse": "{WIDE_INDENT}🌰 Agent {agent.unique_id} ({agent.life_stage}) is dispersing {n_seeds} seeds..."
+        }
+    },
+    "sim": {
+        "Vegetation": {
+            "on_start": "### 🌵 Simulation started (maximum number of steps: {sim.num_steps}) ###",
+            "on_step": "# 🕰️ Time passes. It is the year {sim.steps} #",
+            "on_manage": "# ⛏ Adding {n_agents} agents to the simulation..."
+        }
+    }
+}

vegetation/config/stages.py

@@ -9,6 +9,9 @@ class LifeStage(IntEnum):
     ADULT = 4
     BREEDING = 5
 
+    def __str__(self):
+        return self.name
+
 
 LIFE_STAGE_RGB_VIZ_MAP = {
     LifeStage.DEAD: (100, 0, 0, 1),