customreward

Creating Custom Rewards

Users can implement novel reward functions tailored to specific research questions.

1. Create File: Add a new Python file in the rewards/predefined/ directory (e.g., rewards/predefined/my_custom_reward.py).

2. Inherit & Implement: Define a class inheriting from rewards.base_reward.BaseReward. Implement the __call__(self, cluster_info, current_tasks, current_time) method containing your custom logic. Optionally store the result in self.last_reward.

   # rewards/predefined/my_custom_reward.py
   from rewards.base_reward import BaseReward
   from rewards.registry_utils import register_reward
   import numpy as np
   
   @register_reward("my_custom")  # Choose a unique name
   class MyCustomReward(BaseReward):
       def __init__(self, custom_param=1.0):
           super().__init__()
           self.custom_param = custom_param
   
       def __call__(self, cluster_info, current_tasks, current_time):
           # Example: Penalize variance in CPU utilization across DCs
           cpu_utils = [
               info["__common__"]["cpu_util_percent"]
               for info in cluster_info["datacenter_infos"].values()
           ]
           util_variance = np.var(cpu_utils) if cpu_utils else 0
           reward = -util_variance * self.custom_param
           self.last_reward = reward
           return reward
   

3. Register: Decorate your class with @register_reward("your_unique_name") (imported from rewards.registry_utils).

4. Import in Registry: Add an import for your new class at the top of rewards/reward_registry.py so that it is discovered at runtime.

   # rewards/reward_registry.py
   # ... other imports ...
   from rewards.predefined.my_custom_reward import MyCustomReward  # Add this line
   

Once implemented and imported, your custom reward is available by name to CompositeReward and get_reward_function, and can be referenced in the components section of your reward_config.yaml.