Feature: Option to set the tracking URI for MLflowCallback.
Author: seanswyi
Pull request: yes

# What does this PR do?

Previously, the MLflowCallback was only able to set MLflow experiments or runs. This PR adds the option to also set the tracking URI.

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

Fixes https://github.com/huggingface/transformers/issues/28961


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case. https://github.com/huggingface/transformers/issues/28961
- [x] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

- Trainer maintainers: @muellerzr @pacman100
- Original author of MLflowCallback: @noise-field (https://github.com/huggingface/transformers/commit/c48b16b8da991ba87ca82ed03e33481fc712fbfa)

<!-- Your PR 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:

- text models: @ArthurZucker and @younesbelkada
- vision models: @amyeroberts
- speech models: @sanchit-gandhi
- graph models: @clefourrier

Library:

- flax: @sanchit-gandhi
- generate: @gante
- pipelines: @Narsil
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @muellerzr and @pacman100

Integrations:

- deepspeed: HF Trainer/Accelerate: @pacman100
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc and @younesbelkada

Documentation: @stevhliu and @MKhalusova

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @sanchit-gandhi
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->


Comments:
Sounds good, thanks for the feedback @amyeroberts! I just made another minor change in the docstring and committed it: the default value of `MLFLOW_TRACKING_URI` should be an empty string rather than `None`.

Is there any way to rerun tests? The failed `tests_torch` seems to be a timeout-related issue and wasn't there before my most recent commit. If passing all tests is absolutely necessary regardless of whether or not my code is related to it, then I feel like a simple rerun may solve this.

```Python
Some tests failed!

============================= FAILURES SHORT STACK =============================
_ LayoutLMTokenizationTest.test_add_tokens [type (microsoft/layoutlm-base-uncased)] _


self = <huggingface_hub.utils._http.UniqueRequestIdAdapter object at 0x7f75b9f8c730>
request = <PreparedRequest [GET]>, stream = True
timeout = Timeout(connect=10, read=10, total=None), verify = True, cert = None
proxies = OrderedDict([('no', '127.0.0.1,localhost,circleci-internal-outer-build-agent')])

        try:
            resp = conn.urlopen(
                method=request.method,
                url=url,
                body=request.body,
                headers=request.headers,
                redirect=False,
                assert_same_host=False,
                preload_content=False,
                decode_content=False,
                retries=self.max_retries,
                timeout=timeout,
                chunked=chunked,
            )
        ...
        except (_SSLError, _HTTPError) as e:
            if isinstance(e, _SSLError):
                # This branch is for urllib3 versions earlier than v1.22
                raise SSLError(e, request=request)
            elif isinstance(e, ReadTimeoutError):
>               raise ReadTimeout(e, request=request)
E               requests.exceptions.ReadTimeout: (ReadTimeoutError("HTTPSConnectionPool(host='huggingface.co', port=443): Read timed out. (read timeout=10)"), '(Request ID: e8e1debd-e8f7-45eb-b23d-fa63400a0c0b)')

../.pyenv/versions/3.8.12/lib/python3.8/site-packages/requests/adapters.py:532: ReadTimeout




Exited with code exit status 255
```

@seanswyi I can re-run tests for you. I've just set `tests_torch` run off again

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_29032). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@seanswyi As everything is passing and both @muellerzr and I approve, let's just merge and if @noise-field has any comments we can do a follow-up PR. 

Thanks again for adding! 

I'm an MLflow maintainer. Found this PR while I'm investigating https://github.com/mlflow-automation/mlflow/actions/runs/7949278234/job/21702921711#step:15:2462. What if we run this code like this?

```python
import mlflow
import os

assert "MLFLOW_TRACKING_URI" not in os.environ
mlflow.set_tracking_uri("sqlite:///my.db")

# train
...

# Since MLFLOW_TRACKING_URI is not set, the tracking URI is set to an empty string, which is undesired in this case.
```

@harupy Could you elaborate a bit on why you think that would be better? The reason I wrote the initial PR the way I did is because to me it made sense to allow the user to keep `MLFLOW_TRACKING_URI` as an environment variable along with the other MLflow-related environment variables such as `MLFLOW_RUN_NAME`. I also left the default value as an empty string because that seemed to be consistent with MLflow's documentation as well (https://mlflow.org/docs/latest/python_api/mlflow.html#mlflow.set_tracking_uri). I feel like the snippet you provided takes away that control from the user, as it actually seems to be enforcing the user keep the tracking URI to a hardcoded value of `"sqlite:///my.db"`.

Please let me know if I'm misunderstanding or missing something. I haven't been able to fully look into the failed tests you provided but it seems like there's an issue between this PR and the [autologging module](https://github.com/mlflow/mlflow/blob/e27821c256ff08879f17721adb0b034552f2cf2c/mlflow/tracking/fluent.py#L2013)? 

@seanswyi

> I feel like the snippet you provided takes away that control from the user, as it actually seems to be enforcing the user keep the tracking URI to a hardcoded value of "sqlite:///my.db".

True, but there may be users who use the code like my snippet (we do something similar in our tests, no `MLFLOW_TRACKING_URI` environment variable, just call `mlflow.set_tracking_uri` with a temporary sqlite file). Can we call `set_tracking_uri` only when `MLFLOW_TRACKING_URI` exists?


@harupy If you believe that only calling the function when the environment variable exists is better then how about changing the current code to something like this?

```Python
if "MLFLOW_TRACKING_URI" in os.environ:
    self._tracking_uri = os.environ["MLFLOW_TRACKING_URI"]
    logger.debug(f"MLflow tracking URI is set to {self._tracking_uri}")
    self._ml_flow.set_tracking_uri(self._tracking_uri)
else:
    logger.debug(f"Environment variable `MLFLOW_TRACKING_URI` is not provided and therefore will not be explicitly set.")
```

Just curious, is there any reason why a temporary SQLite DB is used? It seems like the documentation would suggest that setting a tracking URI isn't necessary and that the data will just be stored locally at `./mlruns`.

Yeah that looks good to me.

> is there any reason why a temporary SQLite DB is used?

- To clean up runs and experiments after each test invocation.
- To run tests faster.


Hi! I'm here because this PR broke my code :) I think there might be a compatibility break here? 

In my code I call `mlflow.set_tracking_uri(my_mlflow_tracking_uri)`  , but I never set the MLFLOW_TRACKING_URI env var. Is there a way to change this logic so that it checks if the tracking_uri was already set in the mlflow library? In my case it appears that the code 
```
        self._tracking_uri = os.getenv("MLFLOW_TRACKING_URI", "")
        self._experiment_name = os.getenv("MLFLOW_EXPERIMENT_NAME", None)
        self._flatten_params = os.getenv("MLFLOW_FLATTEN_PARAMS", "FALSE").upper() in ENV_VARS_TRUE_VALUES
        self._run_id = os.getenv("MLFLOW_RUN_ID", None)
        logger.debug(
            f"MLflow experiment_name={self._experiment_name}, run_name={args.run_name}, nested={self._nested_run},"
            f" tags={self._nested_run}"
            f" tags={self._nested_run}, tracking_uri={self._tracking_uri}"
        )
```
means that the tracking_uri is getting overriden to be "" because I didn't set the env var.

I can change my implementation to be compatible with the change that was merged, but others may have similar issues if their implementation was similar to mine.

Hi @njbrake! Thanks for the comment, that's something that I missed earlier. The MLflow API seems to provide a [`mlflow.get_tracking_uri()`](https://mlflow.org/docs/latest/python_api/mlflow.html#mlflow.get_tracking_uri) and [`mlflow.is_tracking_uri_set()`](https://mlflow.org/docs/latest/python_api/mlflow.html#mlflow.is_tracking_uri_set) function. It would probably be better to check before initializing `self._tracking_uri`.

Maybe I could incorporate this in the other PR (https://github.com/huggingface/transformers/pull/29096) that addresses @harupy's previous comment as well? (cc. @amyeroberts)

@seanswyi Yes please!