Nested MLflow logging with cross-validation
Author: helena-balabin
Pull request: no

First of all: I apologize for not using the bug/feature request templates, I believe I'm having more of a general question here.

I was wondering if there is any way in which the `MLflowCallback` can be **used in conjunction with a cross-validation training procedure**? 

In each split, I am initializing a new model, and using a new training and test dataset. Also, I am initializing a new `Trainer` with respective `TrainingArguments` in each split. **Ideally I would use a parent run, and log the run of each split as a nested child run.** I have attached the relevant code snippet:

    mlflow.set_tracking_uri(logging_uri_mlflow)
    mlflow.set_experiment('NLP Baseline')

    # Start a parent run so that all CV splits are tracked as nested runs
    mlflow.start_run(run_name='Parent Run')

    for indices in train_test_splits:
        # Initialize tokenizer and model
        tokenizer = AutoTokenizer.from_pretrained(model_type)
        model = AutoModelForSequenceClassification.from_pretrained(model_type, num_labels=len(unique_tags))

        # Encode all text evidences, pad and truncate to max_seq_len
        train_evidences = tokenizer(evidences_text[indices["train_idx"]].tolist(), truncation=True, padding=True)
        test_evidences = tokenizer(evidences_text[indices["test_idx"]].tolist(), truncation=True, padding=True)
        train_labels = labels[indices["train_idx"]].tolist()
        test_labels = labels[indices["test_idx"]].tolist()
        train_dataset = CustomDataset(encodings=train_evidences, labels=train_labels)
        test_dataset = CustomDataset(encodings=test_evidences, labels=test_labels)

        # Note that due to the randomization in the batches, the training/evaluation is slightly
        # different every time
        training_args = TrainingArguments(
            # label_names
            output_dir=output_dir,
            num_train_epochs=epochs,  # total number of training epochs
            logging_steps=100,
            report_to=["mlflow"],  # log via mlflow
            do_train=True,
            do_predict=True,
        )

        # Initialize Trainer based on the training dataset
        trainer = Trainer(
            model=model,
            args=training_args,
            train_dataset=train_dataset,
        )
        # Train
        trainer.train()

        # Make predictions for the test dataset
        predictions = trainer.predict(test_dataset=test_dataset).predictions
        predicted_labels = np.argmax(predictions, axis=1)
        # Use macro average for now
        f1_scores.append(f1_score(test_labels, predicted_labels, average="macro"))

    logger.info(f'Mean f1-score: {np.mean(f1_scores)}')
    logger.info(f'Std f1-score: {np.std(f1_scores)}')

    # End parent run
    mlflow.end_run()

However, this results in the following exception: 
`Exception: Run with UUID d2bf3cf7cc7b4e359f4c4db098604350 is already active. To start a new run, first end the current run with mlflow.end_run(). To start a nested run, call start_run with nested=True`

**I assume that the `nested=True` parameter is required in the `self._ml_flow.start_run()` call in the `setup()` function of `MLflowCallback`?** I tried to remove the MLflowCallback from the Trainer and add a custom callback class that overrides the default `TrainerCallback` in the same way that `MLflowCallback` does, except for using `self._ml_flow.start_run(nested=True)`. Still, that results in separate individual runs being logged, rather than a nested parent run with child runs. 

Are there any best practices for using huggingface models with mlflow logging in a cross-validation procedure? Thanks a lot in advance for any advice or useful comments on that! :smile: 

Comments:
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.