Batch planning

PyAres/Demo/Planners/planner_wiki.py

@@ -1,23 +1,27 @@
-from PyAres import AresPlannerService, PlanRequest, PlanResponse, AresDataType
+from PyAres import *
+from typing import List
 import random
 
-def generate_plan(request: PlanRequest) -> PlanResponse:
-    planned_values = []
-    names = []
+def generate_plans(request: PlanRequest) -> List[Plan]:    
+    plans = []
+    for i in request.previous_plan_status_codes:
+        print(f"Status Code: {i}")
 
-    # Iterate through every parameter configured in the ARES Experiment
-    for param in request.parameters:
-        # Simple Logic: Pick a random value within the allowed range
-        val = random.uniform(param.minimum_value, param.maximum_value)
+    # Make as many plans as was requested
+    for i in range(request.batch_size):
+        current_params = []
+        for param in request.parameters:
+            # Simple Logic: Pick a random value within the allowed range
+            val = random.uniform(param.minimum_value, param.maximum_value)
+            current_params.append(PlannedParameter(param.name, val))
 
-        names.append(param.name)
-        planned_values.append(val)
+        plans.append(Plan(current_params, Outcome.SUCCESS))
 
-    return PlanResponse(parameter_names=names, parameter_values=planned_values)
+    return plans
 
 if __name__ == "__main__":
     service = AresPlannerService(
-        generate_plan, 
+        generate_plans, 
         "Random Search Planner", 
         "This planner picks random values within bounds.", 
         "1.0.0"

PyAres/Models/__init__.py

@@ -1,4 +1,4 @@
-from .ares_data_models import AresDataType, Outcome, RequestMetadata, AresSchemaEntry, Quantity, QuantitySchema, Limits
+from .ares_data_models import AresDataType, Outcome, RequestMetadata, AresSchemaEntry, Quantity, QuantitySchema, Limits, PlanStatusCode
 
 __all__ = [
     "AresDataType",
@@ -7,5 +7,6 @@
     "AresSchemaEntry",
     "Quantity",
     "QuantitySchema",
-    "Limits"
+    "Limits",
+    "PlanStatusCode"
 ]

PyAres/Models/ares_data_models.py

@@ -30,6 +30,12 @@ class Outcome(Enum):
     WARNING = 3
     CANCELED = 4
 
+class PlanStatusCode(Enum):
+    PLAN_STATUS_UNSPECIFIED = 0
+    PLAN_ACCEPTED = 1
+    PLAN_UNACHIEVABLE = 2
+    PLAN_FAILED = 3
+
 class RequestMetadata():
     def __init__(self, proto_metadata: request_metadata_pb2.RequestMetadata):
         self.system_name = proto_metadata.system_name

PyAres/Planning/__init__.py

@@ -1,9 +1,12 @@
-from .planner_models import PlanningParameter, PlanRequest, PlanResponse, ParameterHistoryItem
+from .planner_models import PlanningParameter, PlanRequest, PlanResponse, ParameterHistoryItem, Plan, PlannedParameter
 from .planning_service import AresPlannerService
 
 __all__ = [
     "PlanningParameter",
     "PlanRequest",
     "PlanResponse",
-    "AresPlannerService"
+    "AresPlannerService",
+    "Plan",
+    "PlannedParameter",
+    "ParameterHistoryItem",
 ]

PyAres/Planning/planner_models.py

@@ -1,5 +1,6 @@
 from typing import Dict, Any, List, Sequence, Optional
-from ..Models import Outcome, AresDataType, RequestMetadata
+from ..Models import Outcome, AresDataType, RequestMetadata, PlanStatusCode
+from enum import Enum
 
 class ParameterHistoryItem:
     """ Represents a single historical parameter item """
@@ -28,9 +29,15 @@ class PlanningParameter:
     Designed to provide a more user-friendly abstraction for the user to interact
     with planning parameters through.
     """
-    def __init__(self, name: str, minimum_value: float, 
-                 maximum_value: float, param_history: list[ParameterHistoryItem], data_type: AresDataType, 
-                 is_planned: bool, is_result: bool, planner_name: str, initial_value = None):
+    def __init__(self, name: str, 
+                 data_type: AresDataType,
+                 minimum_value: float = None, 
+                 maximum_value: float = None, 
+                 param_history: list[ParameterHistoryItem] = [],
+                 is_planned: bool = False, 
+                 is_result: bool = False, 
+                 planner_name: str = "", 
+                 initial_value = None):
         """
         Initializes a PlanningParameter.
 
@@ -107,7 +114,13 @@ class PlanRequest:
 
     Designed to provide a more user-friendly abstraction for interacting with a plan request message.
     """
-    def __init__(self, parameters: list[PlanningParameter], settings: Dict[str, Any], analysis_results: Sequence[float], metadata: RequestMetadata = RequestMetadata.from_default_values()):
+    def __init__(self, 
+                parameters: list[PlanningParameter], 
+                settings: Dict[str, Any], 
+                analysis_results: Sequence[float],
+                metadata: RequestMetadata = RequestMetadata.from_default_values(),
+                batch_size: int = 1,
+                previous_plan_status_codes: List[PlanStatusCode] = None):
         """
         Initializes a PlanRequest.
 
@@ -117,8 +130,14 @@ def __init__(self, parameters: list[PlanningParameter], settings: Dict[str, Any]
         self.parameters = parameters
         self.settings = settings
         self.analysis_results = analysis_results
+        self.batch_size = batch_size
         self.request_metadata = metadata
 
+        if previous_plan_status_codes is None:
+            self.previous_plan_status_codes = []
+        else:    
+            self.previous_plan_status_codes = previous_plan_status_codes
+
     def __str__(self) -> str:
         param_str = "\n ".join(self.parameter_names)
         settings_str = "\n ".join([f"{k}: {v}" for k, v in self.settings.items()])
@@ -127,13 +146,15 @@ def __str__(self) -> str:
         metadata_str = str(self.request_metadata).replace('\n', '\n ')
         return (f"PlanRequest object with:\n"
                 f"parameters:\n"
-                f" {param_str}\n"
+                f"{param_str}\n"
                 f"settings:\n"
-                f" {settings_str}\n"
+                f"{settings_str}\n"
                 f"analysis_results:\n"
-                f" {analysis_str}\n"
+                f"{analysis_str}\n"
                 f"request_metadata:\n"
-                f"{metadata_str}")
+                f"{metadata_str}"
+                f"batch_size:\n"
+                f"{self.batch_size}")
 
     def __repr__(self) -> str:
         return self.__str__()
@@ -175,6 +196,8 @@ def __init__(self,
             parameter_names: A list of names associated with planned parameters.
             parameter_values: A list of values associated with planned parameters. 
             parameter_data: A python dictionary of key:value pairs of planned parameters and planned values
+            outcome: An enum of type Outcome that determines whether the planning process succeeded or not, defaults to SUCCESS
+            error_string: An optional string for specifying planning failure reasons to be relayed to ARES
         """
         if parameter_data is not None:
             self.parameter_names = list(parameter_data.keys())
@@ -200,4 +223,39 @@ def __str__(self):
                 f" error_string: {self.error_string}\n")
 
     def __repr__(self) -> str:
-        return self.__str__()
+        return self.__str__()
+
+class PlannedParameter:
+    def __init__(self, parameter_name: str, parameter_value: Any):
+        self.parameter_name = parameter_name
+        self.parameter_value = parameter_value
+
+    def __str__(self):
+        # A clean, readable key-value output
+        return f"{self.parameter_name}: {self.parameter_value}"
+
+    def __repr__(self):
+        # The !r formatting flag automatically wraps strings in quotes and calls __repr__ on the values
+        return f"PlannedParameters(parameter_name={self.parameter_name!r}, parameter_value={self.parameter_value!r})"
+
+class Plan:
+    def __init__(self, planned_parameters: List[PlannedParameter], outcome: Outcome, error_string: str = ""):
+        self.planned_parameters = planned_parameters
+        self.outcome = outcome
+        self.error_string = error_string
+
+    def __str__(self):
+        base_str = f"Plan (Outcome: {self.outcome})"
+
+        # Format the list of parameters into a readable string
+        if self.planned_parameters:
+            params_str = ", ".join(str(p) for p in self.planned_parameters)
+            base_str += f" | Parameters: [{params_str}]"
+
+        if self.error_string:
+            base_str += f" - Error: '{self.error_string}'"
+
+        return base_str
+
+    def __repr__(self):
+        return f"Plan(planned_parameters={self.planned_parameters!r}, outcome={self.outcome!r}, error_string={self.error_string!r})"

PyAres/Planning/planning_service.py

@@ -21,13 +21,15 @@
 from ..Utils import ares_data_type_utils
 from ..Utils import ares_struct_utils
 from ..Utils import ares_outcome_utils
+from ..Utils import ares_plan_status_code_utils
+from ..Utils import plan_response_utils
 
 # Import python models
 from ..Models import ares_data_models, Limits
 from .planner_models import *
 
 # Type hint for the user's custom planning logic
-PlanLogicFunction = Callable[[PlanRequest], Union[PlanResponse, Awaitable[PlanResponse]]]
+PlanLogicFunction = Callable[[PlanRequest], Union[PlanResponse, Awaitable[PlanResponse], List[Plan], Awaitable[List[Plan]]]]
 
 class AresPlannerServiceWrapper(planner_service_grpc.AresRemotePlannerServiceServicer):
     """
@@ -122,7 +124,9 @@ def Plan(self, request: plan_pb2.PlanningRequest, context) -> plan_pb2.PlanningR
         python_request = PlanRequest(parameters=parameters, 
                                      settings=ares_struct_utils.ares_struct_to_dict(request.adapter_settings), 
                                      analysis_results=list(request.analysis_results),
-                                     metadata=RequestMetadata(request.metadata))
+                                     metadata=RequestMetadata(request.metadata),
+                                     batch_size=request.batch_size,
+                                     previous_plan_status_codes=[ares_plan_status_code_utils.proto_plan_status_to_python_plan_status(c) for c in request.previous_plan_status_codes])
 
         #Handle call using the user's custom planning logic 
         response_proto = plan_pb2.PlanningResponse()
@@ -139,20 +143,23 @@ def Plan(self, request: plan_pb2.PlanningRequest, context) -> plan_pb2.PlanningR
             response_proto.planning_outcome = ares_outcome_enum_pb2.FAILURE
             return response_proto
 
-        if not isinstance(python_response, PlanResponse):
-            response_proto.error_string = "The returned response from the user planning method was not a plan response, and thus was invalid."
-            response_proto.planning_outcome = ares_outcome_enum_pb2.FAILURE
-            return response_proto
-
-        response_proto.planning_outcome = ares_outcome_utils.python_ares_outcome_to_proto_ares_outcome(python_response.outcome)
-        response_proto.error_string = python_response.error_string
+        if isinstance(python_response, PlanResponse):
+            response_proto.planning_outcome = ares_outcome_utils.python_ares_outcome_to_proto_ares_outcome(python_response.outcome)
+            response_proto.error_string = python_response.error_string
+
+            for i in range(len(python_response.parameter_names)):
+                planned_parameter = plan_pb2.PlannedParameter(parameter_value=ares_value_utils.create_ares_value(python_response.parameter_values[i]))
+                planned_parameter.parameter_name = python_response.parameter_names[i]
+                new_planned_parameter = response_proto.planned_parameters.add()
+                new_planned_parameter.CopyFrom(planned_parameter)
 
-        for i in range(len(python_response.parameter_names)):
-            planned_parameter = plan_pb2.PlannedParameter(parameter_value=ares_value_utils.create_ares_value(python_response.parameter_values[i]))
-            planned_parameter.parameter_name = python_response.parameter_names[i]
-            new_planned_parameter = response_proto.planned_parameters.add()
-            new_planned_parameter.CopyFrom(planned_parameter)
+        elif isinstance(python_response, List) and all(isinstance(item, Plan) for item in python_response):
+            response_proto.plans.extend(plan_response_utils.python_plan_to_proto_plan(p) for p in python_response)
 
+        else:
+            response_proto.error_string = "The returned response from the user planning method was not valid, users must return either a list of plans or a plan response."
+            response_proto.planning_outcome = ares_outcome_enum_pb2.FAILURE
+
         print("Sending Plan Response.....")
         return response_proto
 

PyAres/Utils/ares_outcome_utils.py

@@ -2,10 +2,9 @@
 from ares_datamodel import ares_outcome_enum_pb2
 from typing import cast
 
-def python_ares_outcome_to_proto_ares_outcome(py_value: Outcome) -> ares_outcome_enum_pb2.Outcome:
+def python_ares_outcome_to_proto_ares_outcome(py_value: Outcome) -> int:
   """ A method to convert from the python AresDataType class to the protobuf version """
-  val = cast(ares_outcome_enum_pb2.Outcome, py_value.value)
-  return val
+  return py_value.value
 
 def proto_ares_outcome_to_python_ares_outcome(proto_value: ares_outcome_enum_pb2.Outcome) -> Outcome:
   """ A method to convert from the protobuf AresDataType class to the python version """

PyAres/Utils/ares_plan_status_code_utils.py

@@ -0,0 +1,9 @@
+from ..Models import PlanStatusCode
+from ares_datamodel.planning import plan_pb2
+from typing import cast
+
+def python_plan_status_to_proto_plan_status(py_value: PlanStatusCode) -> int:
+    return py_value.value
+
+def proto_plan_status_to_python_plan_status(proto_value: plan_pb2.PlanStatusCode) -> PlanStatusCode:
+  return PlanStatusCode(proto_value)

PyAres/Utils/ares_value_utils.py

@@ -7,7 +7,7 @@
 from . import ares_data_type_utils
 from ..Models import AresDataType, Quantity
 
-def ares_value_to_py(ares_value: ares_struct_pb2.AresValue):
+def ares_value_to_py(ares_value: ares_struct_pb2.AresValue) -> Any:
     """Converts an AresValue protobuf message to a Python native type."""
     field = ares_value.WhichOneof("kind")
     if field == "null_value":

PyAres/Utils/device_status_code_utils.py

@@ -2,10 +2,9 @@
 from ares_datamodel import command_status_code_pb2
 from typing import cast
 
-def python_status_code_to_proto_status_code(py_value: StatusCode) -> command_status_code_pb2.CommandStatusCode:
+def python_status_code_to_proto_status_code(py_value: StatusCode) -> int:
   """ A method to convert from the python StatusCode enum class to the protobuf version """
-  val = cast( command_status_code_pb2.CommandStatusCode, py_value.value)
-  return val
+  return py_value.value
 
 def proto_status_code_to_python_status_code(proto_value:  command_status_code_pb2.CommandStatusCode) -> StatusCode:
   """ A method to convert from the protobuf StatusCode enum class to the python version """

PyAres/Utils/plan_response_utils.py

@@ -1,6 +1,6 @@
 from ares_datamodel.planning import plan_pb2
-from ..Planning import PlanResponse
-from . import ares_outcome_utils
+from ..Planning import PlanResponse, Plan
+from . import ares_outcome_utils, planning_param_utils
 
 def proto_plan_response_to_python(proto_response: plan_pb2.PlanningResponse) -> PlanResponse:
   param_names = []
@@ -14,3 +14,11 @@ def proto_plan_response_to_python(proto_response: plan_pb2.PlanningResponse) ->
   python_response.outcome = ares_outcome_utils.proto_ares_outcome_to_python_ares_outcome(proto_response.planning_outcome)
   python_response.error_string = proto_response.error_string
   return python_response
+
+def python_plan_to_proto_plan(python_plan: Plan) -> plan_pb2.Plan:
+  proto_plan = plan_pb2.Plan()
+  proto_plan.planned_parameters.extend(planning_param_utils.convert_python_planned_param_to_proto_planned_param(p) for p in python_plan.planned_parameters)
+  proto_plan.planning_outcome = ares_outcome_utils.python_ares_outcome_to_proto_ares_outcome(python_plan.outcome)
+  proto_plan.error_string = python_plan.error_string
+
+  return proto_plan

PyAres/Utils/planning_param_utils.py

@@ -1,6 +1,5 @@
-from ..Planning import PlanningParameter, PlanResponse
+from ..Planning import PlanningParameter, PlannedParameter
 from ares_datamodel.planning import plan_pb2
-from typing import List
 from . import param_history_info_utils
 from . import ares_data_type_utils
 from . import ares_value_utils
@@ -17,4 +16,29 @@ def convert_python_plan_param_to_proto(param: PlanningParameter) -> plan_pb2.Pla
   new_proto_param.planner_name = param.planner_name
   ares_value_utils.py_to_ares_value(param.initial_value, new_proto_param.initial_value)
 
-  return new_proto_param 
+  return new_proto_param 
+
+def convert_proto_plan_param_to_python(param: plan_pb2.PlanningParameter) -> PlanningParameter:
+  new_python_param = PlanningParameter(param.parameter_name, ares_data_type_utils.proto_ares_type_to_python_ares_type(param.data_type))
+  new_python_param.minimum_value = param.minimum_value
+  new_python_param.maximum_value = param.maximum_value
+  new_python_param.param_history.extend([param_history_info_utils.convert_proto_param_history_to_python(p) for p in param.parameter_history])
+  new_python_param.is_planned = param.is_planned
+  new_python_param.is_result = param.is_result
+  new_python_param.planner_name = param.planner_name
+
+  if(param.initial_value is not None):
+    new_python_param.initial_value = ares_value_utils.ares_value_to_py(param.initial_value)
+
+  return new_python_param
+
+def convert_python_planned_param_to_proto_planned_param(param: PlannedParameter) -> plan_pb2.PlannedParameter:
+  new_proto_param = plan_pb2.PlannedParameter()
+  new_proto_param.parameter_name = param.parameter_name
+  ares_value_utils.py_to_ares_value(param.parameter_value, new_proto_param.parameter_value)
+
+  return new_proto_param
+
+def convert_proto_planned_param_to_proto_planned_param(param: plan_pb2.PlannedParameter) -> PlannedParameter:
+  new_python_param = PlannedParameter(parameter_name=param.parameter_name, parameter_value=ares_value_utils.ares_value_to_py(param.parameter_value))
+  return new_python_param