more prints

mosaicml · Nov 9, 2023 · e8274a1 · e8274a1
1 parent d47cf4d
commit e8274a1
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 7 deletions.
diff --git a/llmfoundry/utils/builders.py b/llmfoundry/utils/builders.py
@@ -118,9 +118,12 @@ def build_callback(name: str, kwargs: Dict[str, Any]) -> Callback:
         return EarlyStopper(**kwargs)
     elif name == 'hf_checkpointer':
         mlflow_logging_config = kwargs.pop('mlflow_logging_config', None)
+        print(f"build_callback::mlflow_logging_config={mlflow_logging_config}")
+        print(f"build_callback::isinstance(mlflow_logging_config, DictConfig)={isinstance(mlflow_logging_config, DictConfig)}")
         if isinstance(mlflow_logging_config, DictConfig):
             mlflow_logging_config = om.to_object(mlflow_logging_config)
         kwargs['mlflow_logging_config'] = mlflow_logging_config
+        print(f"build_callback::kwargs['mlflow_logging_config']={kwargs['mlflow_logging_config']}")
         return HuggingFaceCheckpointer(**kwargs)
     else:
         raise ValueError(f'Not sure how to build callback: {name}')

diff --git a/tests/test_hf_conversion_script.py b/tests/test_hf_conversion_script.py
@@ -422,13 +422,13 @@ def test_huggingface_conversion_callback(model: str, tmp_path: pathlib.Path,
 
     if dist.get_global_rank() == 0:
         if log_to_mlflow:
-            # mlflow_logger_mock.save_model.assert_called_with(
-            #     flavor='transformers',
-            #     transformers_model=ANY,
-            #     path=ANY,
-            #     task='text-generation',
-            #     metatdata={'task': 'llm/v1/completions'}
-            # )
+            mlflow_logger_mock.save_model.assert_called_with(
+                flavor='transformers',
+                transformers_model=ANY,
+                path=ANY,
+                task='text-generation',
+                metatdata={'task': 'llm/v1/completions'}
+            )
             assert mlflow_logger_mock.register_model.call_count == 1
         else:
             assert mlflow_logger_mock.save_model.call_count == 0