connect the existing mlflow run id

mosaicml · Oct 26, 2024 · bb0dd6a · bb0dd6a
1 parent 6c5fb05
commit bb0dd6a
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/llmfoundry/callbacks/hf_checkpointer.py b/llmfoundry/callbacks/hf_checkpointer.py
@@ -150,6 +150,9 @@ def _log_model_multiprocess(
     log.info("----------------- REACHED MLFLOW LOG MODEL -----------------")
     # monkey patch to prevent duplicate tokenizer upload
     import mlflow
+    mlflow.start_run(
+        run_id=mlflow_logger._run_id,
+    )
     original_save_model = mlflow.transformers.save_model
     def save_model_patch(*args: Any, **kwargs: Any):
         original_save_model(*args, **kwargs)