Repository: mlflow/mlflow

Issue:
[BUG] Postgresql sqlalchemy.exc.DataError is not caught causing the server to hang for over 60 seconds
Thank you for submitting an issue. Please refer to our [issue policy](https://www.github.com/mlflow/mlflow/blob/master/ISSUE_POLICY.md) for additional information about bug reports. For help with debugging your code, please refer to [Stack Overflow](https://stackoverflow.com/questions/tagged/mlflow).

**Please fill in this bug report template to ensure a timely and thorough response.**

### Willingness to contribute
The MLflow Community encourages bug fix contributions. Would you or another member of your organization be willing to contribute a fix for this bug to the MLflow code base?

- [ ] Yes. I can contribute a fix for this bug independently.
- [ ] Yes. I would be willing to contribute a fix for this bug with guidance from the MLflow community.
- [x] No. I cannot contribute a bug fix at this time.

### System information
- **Have I written custom code (as opposed to using a stock example script provided in MLflow)**:
- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: CentOS Linux 7
- **MLflow installed from (source or binary)**: binary
- **MLflow version (run ``mlflow --version``)**: mlflow, version 1.23.1
- **Python version**: Python 3.9.5
- **npm version, if running the dev UI**: N/A
- **Exact command to reproduce**: 

1. Using the backend postgresql+psycopg2 (this error does not occur with sqllite)
2. Run this python method with a non-string/int experiment_id:
```
# this gives an unexpected error and makes the server unresponsive for over 60 seconds
mlflow.start_run(experiment_id={"test":"test"}, run_name="Hello World")
```

### Describe the problem
Describe the problem clearly here. Include descriptions of the expected behavior and the actual behavior.

When I run `mlflow.start_run(experiment_id={"test":"test"}, run_name="Hello World")` where the experiment id is a non-string or integer, instead of getting an expected input validation error, the request hangs for over 60 seconds effectively making the server go unresponsive, and finally returns the error: 
```
MlflowException: API request to http://localhost:5000/api/2.0/mlflow/runs/create failed with exception HTTPConnectionPool(host='localhost', port=5000): Max retries exceeded with url: /api/2.0/mlflow/runs/create (Caused by ResponseError('too many 500 error responses'))
```

This does not happen using the sqllite backend. 

Server Error: 
```
sqlalchemy.exc.DataError: (psycopg2.errors.InvalidTextRepresentation) invalid input syntax for type integer: "{'test': 'test'}"
```
### Tracking information

<!-- This section is optional -->

For bugs related to the tracking features (e.g. mlflow should log a run in my database but it doesn't), please insert the following code in your python script / notebook where you encountered the bug and run it:

```python
print("MLflow version:", mlflow.__version__)
print("Tracking URI:", mlflow.get_tracking_uri())
print("Artifact URI:", mlflow.get_artifact_uri())
```

Then, make sure the printed out information matches what you expect and paste it (with sensitive information masked) in the box below:

```
MLflow version: 1.25.1
Tracking URI: http://localhost:5000/
Artifact URI: local-artifacts/0/68e0c4bee1a14515b03884c56e221f9b/artifacts
```

If you know the command that was used to launch your tracking server (e.g. `mlflow server -h 0.0.0.0 -p 5000`), please provide it:

```
mlflow server --backend-store-uri postgresql+psycopg2:user:secret@localhost/mlflow --default-artifact-root=local-artifacts --host 0.0.0.0

```

### Code to reproduce issue
Provide a reproducible test case that is the bare minimum necessary to generate the problem.
```
# this breaks and makes the server unresponsive for over 60 seconds
mlflow.start_run(experiment_id={"test":"test"}, run_name="Hello World")
```

### Other info / logs
Include any logs or source code that would be helpful to diagnose the problem. If including tracebacks, please include the full traceback. Large logs and files should be attached.


### What component(s), interfaces, languages, and integrations does this bug affect?
Components 
- [ ] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code
- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [ ] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: MLflow Model server, model deployment tools, Spark UDFs
- [X] `area/server-infra`: MLflow Tracking server backend
- [X] `area/tracking`: Tracking Service, tracking client APIs, autologging

Interface 
- [ ] `area/uiux`: Front-end, user experience, plotting, JavaScript, JavaScript dev server
- [ ] `area/docker`: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
- [ ] `area/sqlalchemy`: Use of SQLAlchemy in the Tracking Service or Model Registry
- [ ] `area/windows`: Windows support

Language 
- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

Integrations
- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations



Discussion:
I think we need to validate the "experiment_id" argument type (only support string / int type) in `mlflow.start_run`. 
is there any other more limitations for "experiment_id" ? @dbczumar 
thanks :)


Yep, I think that's all that is needed as well, thanks.
Yep! Sounds great to me!

Resolving patch:
diff --git a/mlflow/tracking/fluent.py b/mlflow/tracking/fluent.py
--- a/mlflow/tracking/fluent.py
+++ b/mlflow/tracking/fluent.py
@@ -38,7 +38,7 @@
     MLFLOW_RUN_NAME,
     MLFLOW_RUN_NOTE,
 )
-from mlflow.utils.validation import _validate_run_id
+from mlflow.utils.validation import _validate_run_id, _validate_experiment_id_type
 
 if TYPE_CHECKING:
     import pandas  # pylint: disable=unused-import
@@ -231,6 +231,7 @@ def start_run(
         0  78b3b0d264b44cd29e8dc389749bb4be          yes           CHILD_RUN
     """
     global _active_run_stack
+    _validate_experiment_id_type(experiment_id)
     # back compat for int experiment_id
     experiment_id = str(experiment_id) if isinstance(experiment_id, int) else experiment_id
     if len(_active_run_stack) > 0 and not nested:
diff --git a/mlflow/utils/validation.py b/mlflow/utils/validation.py
--- a/mlflow/utils/validation.py
+++ b/mlflow/utils/validation.py
@@ -345,6 +345,19 @@ def _validate_experiment_name(experiment_name):
         )
 
 
+def _validate_experiment_id_type(experiment_id):
+    """
+    Check that a user-provided experiment_id is either a string, int, or None and raise an
+    exception if it isn't.
+    """
+    if experiment_id is not None and not isinstance(experiment_id, (str, int)):
+        raise MlflowException(
+            f"Invalid experiment id: {experiment_id} of type {type(experiment_id)}. "
+            "Must be one of str, int, or None.",
+            error_code=INVALID_PARAMETER_VALUE,
+        )
+
+
 def _validate_model_name(model_name):
     if model_name is None or model_name == "":
         raise MlflowException("Registered model name cannot be empty.", INVALID_PARAMETER_VALUE)


Test patch:
diff --git a/tests/tracking/fluent/test_fluent.py b/tests/tracking/fluent/test_fluent.py
--- a/tests/tracking/fluent/test_fluent.py
+++ b/tests/tracking/fluent/test_fluent.py
@@ -487,6 +487,14 @@ def test_start_run_defaults_databricks_notebook(
         assert is_from_run(active_run, MlflowClient.create_run.return_value)
 
 
+@pytest.mark.parametrize(
+    "experiment_id", [("a", "b"), {"a", "b"}, ["a", "b"], {"a": 1}, [], (), {}]
+)
+def test_start_run_raises_invalid_experiment_id(experiment_id):
+    with pytest.raises(MlflowException, match="Invalid experiment id: "):
+        start_run(experiment_id=experiment_id)
+
+
 @pytest.mark.usefixtures(empty_active_run_stack.__name__)
 def test_start_run_creates_new_run_with_user_specified_tags():
     mock_experiment_id = mock.Mock()
@@ -553,7 +561,7 @@ def test_start_run_resumes_existing_run_and_sets_user_specified_tags():
 
 def test_start_run_with_parent():
     parent_run = mock.Mock()
-    mock_experiment_id = mock.Mock()
+    mock_experiment_id = "123456"
     mock_source_name = mock.Mock()
 
     active_run_stack_patch = mock.patch("mlflow.tracking.fluent._active_run_stack", [parent_run])