Don't wait for mlflow.log_artifact in Trainer api
Author: avramdj
Pull request: no

### Feature request

It would be ideal if there was an option to make `transformers.Trainer.train` with `HF_MLFLOW_LOG_ARTIFACTS=1` continue training while a separate process/thread logs the artifacts, so the whole run doesn't hang because of a network bottleneck.

I would be happy to discuss the possibilities and/or limitations of something like this (for example if it conflicts with `save_total_limit`).

### Motivation

I've had multiple situations where model artifact logging was almost as slow as training itself. This is especially a problem when training on expensive cloud GPU nodes for reasons I don't even need to explain.

### Your contribution

I would be willing to discuss and potentially contribute to a feature like this once we've discussed it.

Comments:
I think this makes a lot of sense. 

Note that:

> Integrations with reporting platforms are entirely maintained by the developers of those integrations or the community.

Would you like to open a PR 🤗 ?

Hi @ydshieh, thank's for the response. 

Before opening a PR, I'd like to discuss the feasability of such a feature in the first place. 
I can't tell if this would require an aggressive refactor of platform integration callbacks, if so, maybe it isn't worth it at the moment. 

As I mentioned, I have no idea how introducing this concurrency will affect the rest of the Trainer api, it could cause a race condition when there's stuff like `save_total_limit`. Anyway, here's an outline I thought of so far, haven't tested it yet:

```py
class AsyncTrainerCallback(TrainerCallback):
    # ...
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._log_queue = queue.Queue()
        self._worker_thread = threading.Thread(target=self._worker_loop)
        self._shutdown = False
        self._worker_loop.start()

    def _worker_loop(self):
        while not self._shutdown or not self._log_queue.empty():
            task = self._log_queue.get()
            task()
            self._log_queue.task_done()
    
    def _stop_worker(self):
        self._shutdown = True
        if not self._log_queue.empty():
            print("Waiting for logging to finish...")
        self._log_queue.join()
        self._worker_thread.join()
    

class MLflowCallback(AsyncTrainerCallback):
    def on_save(self, args, state, control, **kwargs):
        if self._initialized and state.is_world_process_zero and self._log_artifacts:
            ckpt_dir = f"checkpoint-{state.global_step}"
            artifact_path = os.path.join(args.output_dir, ckpt_dir)

            ### instead of:
            # logger.info(f"Logging checkpoint artifacts in {ckpt_dir}. This may take time.")
            # self._ml_flow.pyfunc.log_model(
            #     ckpt_dir,
            #     artifacts={"model_path": artifact_path},
            #     python_model=self._ml_flow.pyfunc.PythonModel(),
            # )

            ### do:
            task = lambda: self._ml_flow.pyfunc.log_model(
                ckpt_dir,
                artifacts={"model_path": artifact_path},
                python_model=self._ml_flow.pyfunc.PythonModel(),
            )
            self._log_queue.put(task)

    def on_train_end(self, args, state, control, **kwargs):
        # ...
        self._stop_worker()
```

I think it's probably better not to create the worker from the start. Just create the queue and worker inside `on_save`, then put the task to the queue.

WE shouldn't have to worry about `save_total_limit`, as I  don't see it is used in the current `MLflowCallback`, which means it is controlled by the `Trainer` class.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.