Ensemble documentation update (#322)

Expand documentation for Ensembles

[ committed by @billschereriii ]
[ reviewed by @mellis13 @ankona ]

doc/api/smartsim_api.rst

@@ -450,6 +450,8 @@ Model
 Ensemble
 ========
 
+.. _ensemble_api:
+
 .. currentmodule:: smartsim.entity.ensemble
 
 .. autosummary::

doc/experiment.rst

@@ -259,4 +259,69 @@ the following:
 
   cd /scratch/smartsim/Training-Run ; /usr/bin/srun --output /scratch/smartsim/Training-Run/Training-Ensemble_3.out --error /scratch/smartsim/Training-Ensemble_3.err --job-name Training-Ensemble_3-CHTN0UI2TRE7 --nodes=1 --ntasks=24 /scratch/pyenvs/smartsim/bin/python ./train-model.py --LR=0.35000000000000003 &
 
-  wait
+  wait
+
+Prefixing Keys in the Orchestrator
+----------------------------------
+
+If each of multiple ensemble members attempt to use the same code to access their respective models
+in the Orchestrator, the keys by which they do this will overlap and they can end up accessing each
+others' data inadvertently. To prevent this situation, the SmartSim Entity object supports key
+prefixing, which automatically prepends the name of the model to the keys by which it is accessed.
+With this enabled, key overlapping is no longer an issue and ensemble members can use the same code.
+
+Under the hood, calling ensemble.enable_key_prefixing() causes the SSKEYOUT environment variable to
+be set, which in turn causes all keys generated by an ensemble member to be prefixed with its model
+name. Similarly, if the model for the ensemble member has incoming entities (such as those set via
+model.register_incoming_entity() or ensemble.register_incoming_entity()), the SSKEYIN environment
+variable will be set and the keys associated with those inputs will be automatically prefixed. Note
+that entities must register themselves as this is not done by default.
+
+Finally, please note that while prefixing is enabled by default for tensors, datasets, and aggregated
+lists of datasets, a SmartRedis client must manually call Client.use_model_ensemble_prefix() to
+ensure that prefixes are used with models and scripts.
+
+We modify the example above to enable key prefixing as follows:
+
+.. code-block:: bash
+
+  import numpy as np
+  from smartsim import Experiment
+
+  exp = Experiment("Training-Run", launcher="slurm")
+
+  # setup ensemble parameter space
+  learning_rate = list(np.linspace(.01, .5))
+  train_params = {"LR": learning_rate}
+
+  # define resources for all ensemble members
+  sbatch = exp.create_batch_settings(nodes=4,
+                                    time="01:00:00",
+                                    account="12345-Cray",
+                                    queue="gpu")
+
+  # define how each member should run
+  srun = exp.create_run_settings(exe="python",
+                                exe_args="./train-model.py")
+  srun.set_nodes(1)
+  srun.set_tasks(24)
+
+  ensemble = exp.create_ensemble("Training-Ensemble",
+                                params=train_params,
+                                params_as_args=["LR"],
+                                batch_settings=sbatch,
+                                run_settings=srun,
+                                perm_strategy="random",
+                                n_models=4)
+
+  # Enable key prefixing -- note that this should be done
+  # before starting the experiment
+  ensemble.enable_key_prefixing()
+
+  exp.start(ensemble, summary=True)
+
+
+Further Information
+-------------------
+
+For more informtion about Ensembles, please refer to the :ref:`Ensemble API documentation <ensemble_api>`.

smartsim/entity/ensemble.py

@@ -227,14 +227,14 @@ def register_incoming_entity(self, incoming_entity: SmartSimEntity) -> None:
             model.register_incoming_entity(incoming_entity)
 
     def enable_key_prefixing(self) -> None:
-        """If called, all models within this ensemble will prefix their keys with its
+        """If called, each model within this ensemble will prefix its key with its
         own model name.
         """
         for model in self.models:
             model.enable_key_prefixing()
 
     def query_key_prefixing(self) -> bool:
-        """Inquire as to whether each model within the ensemble will prefix its keys
+        """Inquire as to whether each model within the ensemble will prefix their keys
 
         :returns: True if all models have key prefixing enabled, False otherwise
         :rtype: bool

tests/test_ensemble.py

@@ -52,13 +52,13 @@ def step_values(param_names, param_values, n_models = 0):
     return permutations
 
 
-# bad permuation strategy that doesnt return
+# bad permutation strategy that doesn't return
 # a list of dictionaries
 def bad_strategy(names, values, n_models = 0):
     return -1
 
 
-# test bad perm strat that returns a list but of lists
+# test bad perm strategy that returns a list but of lists
 # not dictionaries
 def bad_strategy_2(names, values, n_models = 0):
     return [values]