small testing fixes

mosaicml · Nov 9, 2023 · ba871d7 · ba871d7
1 parent 9819456
commit ba871d7
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 6 deletions.
diff --git a/llmfoundry/callbacks/async_eval_callback.py b/llmfoundry/callbacks/async_eval_callback.py
@@ -10,7 +10,8 @@
                                               RUN_NAME_ENV_VAR)
 from composer.utils import dist
 from composer.utils.misc import create_interval_scheduler
-from mcli.api.runs import ComputeConfig  # TODO: should be available in root
+from mcli.api.runs import \
+    ComputeConfig  # TODO: available in root in mcli 0.5.27+
 
 from mcli import Run, RunConfig, create_run, get_run
 
@@ -29,7 +30,7 @@
 OPTIONAL_PARAMS_FOR_EVAL = {
     'dist_timeout',
     'eval_gauntlet',
-    'fsdp_config',  # fsdp_dict_cfg
+    'fsdp_config',
     'icl_subset_num_batches',
     'loggers',
     'precision',
@@ -40,7 +41,7 @@
 
 def get_run_name(previous_run_name: str, count: int) -> str:
     *name_without_uuid_suffix, _ = previous_run_name.split('-')
-    name_suffix = '-'.join(name_without_uuid_suffix)[:MAX_RUN_NAME_LENGTH]
+    name_suffix = ('-'.join(name_without_uuid_suffix))[:MAX_RUN_NAME_LENGTH]
     return f'eval{count}-{name_suffix}'
 
 

diff --git a/tests/callbacks/test_async_eval_callback.py b/tests/callbacks/test_async_eval_callback.py
@@ -16,7 +16,7 @@ def test_get_run_name():
     a = get_run_name('foo-1234', 0)
     assert a == 'eval0-foo'
 
-    b = get_run_name(50 * 'foo-1234', 1)
+    b = get_run_name(50 * 'foo' + '-1234', 1)
     assert b == 'eval1-foofoofoofoofoofoofoofoofoofoofoofoofoof'
 
 
@@ -149,7 +149,7 @@ def test_get_eval_parameters():
             },
         }],
         'eval_gauntlet': 'eval_gauntlet_example',
-        'fsdp_dict_cfg': {
+        'fsdp_config': {
             'fsdp_cfg_example': 'fsdp_cfg_example'
         },
         'icl_subset_num_batches': 4,
@@ -222,5 +222,16 @@ def test_async_eval_callback_minimal(mock_create_run: MagicMock,
     assert parameters['icl_tasks'] == 'icl_task_example'
     assert parameters['max_seq_len'] == 3
     assert parameters['load_path'] == 'save_folder_example/latest-rank0.pt'
-    assert parameters['models'] == ['model_example']
+    assert parameters['models'] == [{
+        'model_name': 'model_example',
+        'model': {
+            'name': 'model_example',
+            'attn_config': {
+                'foo': 'bar'
+            }
+        },
+        'tokenizer': {
+            'tokenizer_example': 'tokenizer_example'
+        }
+    }]
     assert parameters['run_name'] == 'eval0-foo_bar'  # original run