MlFlow log artefacts
## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: 4.4.2
- Platform: Darwin-20.3.0-x86_64-i386-64bit
- Python version: 3.7.4
- PyTorch version (GPU?): 1.3.1 (False)
- Tensorflow version (GPU?): not installed (NA)
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: No

### Who can help

@sgugger 
<!-- Your issue will be replied to more quickly if you can figure out the right person to tag with @
 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- albert, bert, xlm: @LysandreJik
- blenderbot, bart, marian, pegasus, encoderdecoder,  t5: @patrickvonplaten, @patil-suraj
- longformer, reformer, transfoxl, xlnet: @patrickvonplaten
- fsmt: @stas00
- funnel: @sgugger
- gpt2: @patrickvonplaten, @LysandreJik
- rag: @patrickvonplaten, @lhoestq
- tensorflow: @LysandreJik

Library:

- benchmarks: @patrickvonplaten
- deepspeed: @stas00
- ray/raytune: @richardliaw, @amogkam
- text generation: @patrickvonplaten
- tokenizers: @LysandreJik
- trainer: @sgugger
- pipelines: @LysandreJik

Documentation: @sgugger

Model hub:

- for issues with a model report at https://discuss.huggingface.co/ and tag the model's creator.

HF projects:

- nlp datasets: [different repo](https://github.com/huggingface/nlp)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Examples:

- maintained examples (not research project or legacy): @sgugger, @patil-suraj
- research_projects/bert-loses-patience: @JetRunner
- research_projects/distillation: @VictorSanh

 -->

## Information

Model I am using (Bert, XLNet ...): Bert

The problem arises when using:
* [x] the official example scripts: (give details below)
* [ ] my own modified scripts: (give details below)

The tasks I am working on is:
* [x] an official GLUE/SQUaD task: NER
* [ ] my own task or dataset: (give details below)

## To reproduce

The bug is for the PR #8016.

Steps to reproduce the behavior:

1. MlFlow installed and  the following env variables exported
```
export HF_MLFLOW_LOG_ARTIFACTS=TRUE
export MLFLOW_S3_ENDPOINT_URL=<custom endpont>
export MLFLOW_TRACKING_URI=<custom uri>
export MLFLOW_TRACKING_TOKEN=<custom token>
```
2. Run the token classification example with the following command
```
python run_ner.py \
  --model_name_or_path bert-base-uncased \
  --dataset_name conll2003 \
  --output_dir /tmp/test-ner \
  --do_train \
  --do_eval
```

<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

## Expected behavior

<!-- A clear and concise description of what you would expect to happen. -->

When the training finishes, before the evaluation is performed, the `integrations.MLflowCallback` executes the method `on_train_end`, where if the env variable `HF_MLFLOW_LOG_ARTIFACTS` is set to `TRUE`, it logs the model artifacts to mlflow. 

The problem is, however, when the method `on_train_end` is called and the following line is executed: `self._ml_flow.log_artifacts(args.output_dir)`, the model is not stored on the `args.output_dir`. The model artefacts are stored once the `trainer.save_model()` is called, which is after the training ending. There is no callback in the `trainer.save_model()` that can be called from a `TrainerCallback` to save the model. There is a method `TrainierCallback.on_save()` method, that is called `trainer._maybe_log_save_evaluate()`, but even then the model is not available on the `output_dir`.

Possible solutions would be to extend the `TrainierCallback` with `on_model_save()` callback method, insert the callback in the `trainer.save_model()`. 
Or, a workaround I have now is to change `on_train_end ` with `on_evaluate` in `integrations.MLflowCallback`, that is called after the model is saved in the example script. However, this is not the right solution since it depends on having set the `do_eval` parameter, and it is not semantically correct.



I have no written the `MLFlowCallback` (external integrations are maintained by contributors or the authors of the external libraries themselves) but I can confirm the command will indeed not log the model weights. The callback does get the model in the kwargs, so it's completely possible to get it from there and upload it, like it's done in the `WandbCallback`.<|||||>Thanks @sgugger for your fast reply! 

I tested your suggestion, but it doesn't quite work. I don't know about the `WandbCallback`, but in `MLFlowCallback` in the `on_train_end`, when the model is saved and logged to `mlflow`, the `mlflow` run is ended, and the later loggin of metrics, like evaluation and testing are logged in separate run, which is not what we want.
I don't know if this happens when you create `fake_trainer` with `fake_trainer = Trainer(args=args, model=model, tokenizer=tokenizer)`  or when the artifacts are logged with `self._ml_flow.log_artifacts(temp_dir)`.

Also the current `MLFlowCallback` doesn't log testing metrics. I know this is not a fault of the callback, it is in the `trainer.predict()` method that doesn't have call to `log()` internally. The workaround is to call `trainer.log(metrics)` after 
```
trainer.log_metrics("test", metrics)
trainer.save_metrics("test", metrics)
```
in the example. <|||||>As I said, I haven't written that callback: integrations with reporting platforms are entirely maintained by the developers of those integrations or the community. You can open a PR with your fixes!<|||||>I understand. However, not having a callback hook on the `save_model` would be difficult. 

If somebody is interested, a dirty workaround I did is,
1. Register own `MLflowCallback`

```
trainer = Trainer(
       ...
        callbacks=[MLflowCallback]
)
trainer.remove_callback(transformers.integrations.MLflowCallback)
```
2. Add method in the class:
```
    def log_artifact(self, output_dir):
        if self._initialized:
            logger.info("Logging artifacts. This may take time.")
            self._ml_flow.log_artifacts(output_dir)
```

3. In the `run_ner.py` file, at the very end (or after ` trainer.save_model()`) added
```
ml_flow_callback = trainer.pop_callback(MLflowCallback)
ml_flow_callback.log_artifact(training_args.output_dir)
```
Which removes the `MLflowCallback` and tells to log the model.

I know it is dirty, but if I come up with better solution I will open PR.

Thanks!<|||||>> However, not having a callback hook on the save_model would be difficult.

Not that this hook would be called when each checkpoint is saved, not just at the end of training. So you would not only save the last model.<|||||>You are right, even with my hack of logging the saved model from the `output_dir`, transfers the checkpoint models as well, which is not what we need.

I think modifying `MLflowCallback.on_train_end` with the code from `Trainer._save` should save only the model in temp directory and log it to mlflow. This way, we don't lose the current mlflow run and we dont save everything from the `output_dir`.

```
    def on_train_end(self, args, state, control, model=None, tokenizer=None, **kwargs):
        if self._initialized and state.is_world_process_zero and self._log_artifacts:
            logger.info("Logging artifacts. This may take time.")
            with tempfile.TemporaryDirectory() as temp_dir:

                if not isinstance(model, PreTrainedModel):
                    if isinstance(unwrap_model(model), PreTrainedModel):
                        state_dict = model.state_dict()
                        unwrap_model(model).save_pretrained(temp_dir, state_dict=state_dict)
                    else:
                        logger.info("Trainer.model is not a `PreTrainedModel`, only saving its state dict.")
                        state_dict = model.state_dict()
                        torch.save(state_dict, os.path.join(temp_dir, WEIGHTS_NAME))
                else:
                    state_dict = model.state_dict()
                    model.save_pretrained(temp_dir, state_dict=state_dict)
                if tokenizer is not None:
                    tokenizer.save_pretrained(temp_dir)

                # Good practice: save your training arguments together with the trained model
                torch.save(args, os.path.join(temp_dir, "training_args.bin"))
                self._ml_flow.log_artifacts(temp_dir)
```

If you think this is a good idea, maybe it can be added in the `MLflowCallback` integration.
Thanks!<|||||>That sounds like the best compromise yes.<|||||>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/master/CONTRIBUTING.md) are likely to be ignored.<|||||>Hello, any update on this problem? I am trying to log a model with mlflow but the artifacts aren't registered.

Could you please help me with this ?

Best Regards,<|||||>> Hello, any update on this problem? I am trying to log a model with mlflow but the artifacts aren't registered.
> 
> Could you please help me with this ?
> 
> Best Regards,

Did you export `HF_MLFLOW_LOG_ARTIFACTS` environment variable and set it to `True`?<|||||>I was just trying with `HF_MLFLOW_LOG_ARTIFACTS` set and nothing was appearing in the mlflow artifacts