Add a configurable spot strategy

docs/yaml.md

@@ -76,5 +76,6 @@ Each forge command certain parameters. A yaml file with all the parameters can b
     - Use the `--all` flag to run the script on all the instances in a cluster.
     - E.g. `run_cmd: scripts/run.sh {env} {date} {ip}`
 - **service** - `cluster` or `single`
+- **spot_strategy** - Select the [spot allocation strategy](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/ec2/client/create_fleet.html).
 - **user_data** - Custom script passed to instance. Will be run only once when the instance starts up.
 - **valid_time** - How many hours the fleet will stay up. After this time, all EC2s will be destroyed. The default is 8.

src/forge/__init__.py

@@ -11,7 +11,8 @@
     'destroy_after_failure': True,
     'default_ratio': [8, 8],
     'valid_time': 8,
-    'ec2_max': 768
+    'ec2_max': 768,
+    'spot_strategy': 'price-capacity-optimized'
 }
 
 # Required arguments for each Forge job

src/forge/configure.py

@@ -5,7 +5,7 @@
 import sys
 
 import yaml
-from schema import Schema, And, Optional, SchemaError
+from schema import Schema, And, Optional, Or, SchemaError
 
 from .common import set_config_dir
 
@@ -62,7 +62,15 @@ def check_env_yaml(env_yaml):
         Optional('tags'): And(list, len, error="Invalid AWS tags"),
         Optional('excluded_ec2s'): And(list),
         Optional('additional_config'): And(list),
-        Optional('ec2_max'): And(int)
+        Optional('ec2_max'): And(int),
+        Optional('spot_strategy'): And(str, len,
+                                       Or(
+                                           'lowest-price',
+                                           'diversified',
+                                           'capacity-optimized',
+                                           'capacity-optimized-prioritized',
+                                           'price-capacity-optimized'),
+                                       error='Invalid spot allocation strategy'),
     })
     try:
         validated = schema.validate(env_yaml)

src/forge/create.py

@@ -529,6 +529,7 @@ def _check(x, i):
     gpu = config.get('gpu_flag', False)
     market = config.get('market', DEFAULT_ARG_VALS['market'])
     subnet = config.get('aws_subnet')
+    strategy = config.get('spot_strategy')
 
     market = market[-1] if 'cluster-worker' in n else market[0]
 
@@ -545,7 +546,7 @@ def _check(x, i):
             'AllocationStrategy': 'lowest-price'
         },
         'SpotOptions': {
-            'AllocationStrategy': 'capacity-optimized',
+            'AllocationStrategy': strategy,
             'InstanceInterruptionBehavior': 'terminate',
             'MaintenanceStrategies': {
                 'CapacityRebalance': {

tests/test_main.py

@@ -43,7 +43,7 @@ def _load_admin_cfg(env):
       'home_dir': os.path.dirname(FORGE_DIR), 'yaml_dir': os.path.join(TEST_DIR, 'data'), 'user': 'test_user',
       'ami': 'single_ami', 'forge_version': False, 'config_dir': os.path.join(TEST_CFG_DIR, 'dev'),
       'region': 'us-east-1', 'destroy_after_success': True, 'destroy_after_failure': True, 'default_ratio': [8, 8],
-      'valid_time': 8, 'ec2_max': 768}),
+      'valid_time': 8, 'ec2_max': 768, 'spot_strategy': 'price-capacity-optimized'}),
     # Destroy job; passing the relative path to a yaml
     (['forge', 'destroy', '--yaml', os.path.join(TEST_DIR_REL, 'data', 'single_intermediate.yaml')],
      {'name': 'test-single-intermediate', 'log_level': 'INFO',
@@ -53,7 +53,7 @@ def _load_admin_cfg(env):
       'yaml_dir': os.path.join(TEST_DIR, 'data'), 'user': 'test_user', 'ami': 'single_ami',
       'forge_version': False, 'config_dir': os.path.join(TEST_CFG_DIR, 'dev'), 'region': 'us-east-1',
       'destroy_after_success': True, 'destroy_after_failure': True, 'default_ratio': [8, 8], 'valid_time': 8,
-      'ec2_max': 768}),
+      'ec2_max': 768, 'spot_strategy': 'price-capacity-optimized'}),
     # Destroy job; overriding log_level
     (['forge', 'destroy', '--yaml', os.path.join(TEST_DIR, 'data', 'single_intermediate.yaml'), '--log_level', 'debug'],
      {'name': 'test-single-intermediate', 'log_level': 'DEBUG',
@@ -63,7 +63,7 @@ def _load_admin_cfg(env):
       'home_dir': os.path.dirname(FORGE_DIR), 'yaml_dir': os.path.join(TEST_DIR, 'data'), 'user': 'test_user',
       'ami': 'single_ami', 'forge_version': False, 'config_dir': os.path.join(TEST_CFG_DIR, 'dev'),
       'region': 'us-east-1', 'destroy_after_success': True, 'destroy_after_failure': True, 'default_ratio': [8, 8],
-      'valid_time': 8, 'ec2_max': 768}),
+      'valid_time': 8, 'ec2_max': 768, 'spot_strategy': 'price-capacity-optimized'}),
     # Destroy job; overriding market
     (['forge', 'destroy', '--yaml', os.path.join(TEST_DIR, 'data', 'single_intermediate.yaml'),
       '--market', 'on-demand'],
@@ -73,7 +73,8 @@ def _load_admin_cfg(env):
       'gpu_flag': False, 'app_dir': TEST_DIR, 'src_dir': FORGE_DIR, 'home_dir': os.path.dirname(FORGE_DIR),
       'yaml_dir': os.path.join(TEST_DIR, 'data'), 'user': 'test_user', 'ami': 'single_ami', 'forge_version': False,
       'config_dir': os.path.join(TEST_CFG_DIR, 'dev'), 'region': 'us-east-1', 'destroy_after_success': True,
-      'destroy_after_failure': True, 'default_ratio': [8, 8], 'valid_time': 8, 'ec2_max': 768}),
+      'destroy_after_failure': True, 'default_ratio': [8, 8], 'valid_time': 8, 'ec2_max': 768,
+      'spot_strategy': 'price-capacity-optimized'}),
     # Destroy job; no market
     (['forge', 'destroy', '--yaml', os.path.join(TEST_DIR, 'data', 'single_basic.yaml'), '--forge_env', 'dev'],
      {'name': 'test-single-basic', 'log_level': 'INFO', 'yaml': os.path.join(TEST_DIR, 'data', 'single_basic.yaml'),
@@ -82,7 +83,7 @@ def _load_admin_cfg(env):
       'home_dir': os.path.dirname(FORGE_DIR), 'yaml_dir': os.path.join(TEST_DIR, 'data'), 'user': 'test_user',
       'ami': 'single_ami', 'forge_version': False, 'config_dir': os.path.join(TEST_CFG_DIR, 'dev'),
       'region': 'us-east-1', 'destroy_after_success': True, 'destroy_after_failure': True, 'default_ratio': [8, 8],
-      'valid_time': 8, 'ec2_max': 768}),
+      'valid_time': 8, 'ec2_max': 768, 'spot_strategy': 'price-capacity-optimized'}),
     # Configure job
     (['forge', 'configure'],
      {'forge_version': False, 'job': 'configure', 'log_level': 'INFO'}),
@@ -95,7 +96,7 @@ def _load_admin_cfg(env):
       'home_dir': os.path.dirname(FORGE_DIR), 'yaml_dir': os.path.join(TEST_DIR, 'data'), 'user': 'test_user',
       'ami': 'single_ami', 'forge_version': False, 'config_dir': os.path.join(TEST_CFG_DIR, 'dev'),
       'region': 'us-east-1', 'destroy_after_success': True, 'destroy_after_failure': True, 'default_ratio': [8, 8],
-      'valid_time': 8, 'ec2_max': 768}),
+      'valid_time': 8, 'ec2_max': 768, 'spot_strategy': 'price-capacity-optimized'}),
     # Create job; setting gpu
     (['forge', 'create', '--yaml', os.path.join(TEST_DIR, 'data', 'single_basic.yaml'), '--forge_env', 'dev', '--gpu'],
      {'name': 'test-single-basic', 'log_level': 'INFO', 'yaml': os.path.join(TEST_DIR, 'data', 'single_basic.yaml'),
@@ -104,7 +105,7 @@ def _load_admin_cfg(env):
       'home_dir': os.path.dirname(FORGE_DIR), 'yaml_dir': os.path.join(TEST_DIR, 'data'), 'user': 'test_user',
       'ami': 'single_ami', 'forge_version': False, 'config_dir': os.path.join(TEST_CFG_DIR, 'dev'),
       'region': 'us-east-1', 'destroy_after_success': True, 'destroy_after_failure': True, 'default_ratio': [8, 8],
-      'valid_time': 8, 'ec2_max': 768}),
+      'valid_time': 8, 'ec2_max': 768, 'spot_strategy': 'price-capacity-optimized'}),
 ])
 def test_forge_main(mock_pass, mock_execute, mock_keys, mock_config_dir, cli_call, exp_config, load_admin_cfg):
     """Test the config after calling forge via the command line."""