Repository: mlflow/mlflow

Issue:
[BUG] Azure Databricks Spark UDF not loading model due to tempfile error
### Issues Policy acknowledgement

- [X] I have read and agree to submit bug reports in accordance with the [issues policy](https://www.github.com/mlflow/mlflow/blob/master/ISSUE_POLICY.md)

### Willingness to contribute

No. I cannot contribute a bug fix at this time.

### MLflow version

- Client: 2.7.0
- Tracking URI: databricks

### System information

- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**:  Linux ~20.04.1-Ubuntu SMP
- **Python version**: 3.9.5
- **yarn version, if running the dev UI**:


### Describe the problem

Having trouble loading a saved model using the spark_udf function in Databricks, but only in certain places. When using the spark_udf function directly in a notebook, everything seems OK, but when running it through a function that is "imported" using the magic %run command, I'm getting an error about a tempfile not being able to be created. This seems to be where the error occurs: https://github.com/mlflow/mlflow/blob/master/mlflow/utils/file_utils.py#L830

### Tracking information

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```shell
Python version: 3.9.5
MLflow version: 2.7.0
Tracking URI: databricks
Registry URI: databricks
MLflow environment variables: 
  MLFLOW_CONDA_HOME: /databricks/conda
  MLFLOW_PYTHON_EXECUTABLE: /databricks/spark/scripts/mlflow_python.sh
  MLFLOW_TRACKING_URI: databricks
```


### Code to reproduce issue

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
import mlflow

logged_model_uri = "runs:/id/model"
loaded_model = mlflow.pyfunc.spark_udf(
    spark,
    model_uri=logged_model_uri,
    env_manager="conda",
    result_type="double",
)
```


### Stack trace

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```

File /local_disk0/.ephemeral_nfs/envs/pythonEnv-3ef307dd-1827-4ce9-b993-14ddd54d1d07/lib/python3.9/site-packages/mlflow/pyfunc/__init__.py:1227, in spark_udf(spark, model_uri, result_type, env_manager, params)
   1217     if not sys.platform.startswith("linux"):
   1218         # TODO: support killing mlflow server launched in UDF task when spark job canceled
   1219         #  for non-linux system.
   1220         #  https://stackoverflow.com/questions/53208/how-do-i-automatically-destroy-child-processes-in-windows
   1221         _logger.warning(
   1222             "In order to run inference code in restored python environment, PySpark UDF "
   1223             "processes spawn MLflow Model servers as child processes. Due to system "
   1224             "limitations with handling SIGKILL signals, these MLflow Model server child "
   1225             "processes cannot be cleaned up if the Spark Job is canceled."
   1226         )
-> 1227 pyfunc_backend = get_flavor_backend(
   1228     local_model_path,
   1229     env_manager=env_manager,
   1230     install_mlflow=os.environ.get("MLFLOW_HOME") is not None,
   1231     create_env_root_dir=True,
   1232 )
   1233 if not should_use_spark_to_broadcast_file:
   1234     # Prepare restored environment in driver side if possible.
   1235     # Note: In databricks runtime, because databricks notebook cell output cannot capture
   (...)
   1241     # to wait conda command fail and suddenly get all output printed (included in error
   1242     # message).
   1243     if env_manager != _EnvManager.LOCAL:

File /local_disk0/.ephemeral_nfs/envs/pythonEnv-3ef307dd-1827-4ce9-b993-14ddd54d1d07/lib/python3.9/site-packages/mlflow/models/flavor_backend_registry.py:39, in get_flavor_backend(model_uri, **kwargs)
     37 def get_flavor_backend(model_uri, **kwargs):
     38     if model_uri:
---> 39         with TempDir() as tmp:
     40             if ModelsArtifactRepository.is_models_uri(model_uri):
     41                 underlying_model_uri = ModelsArtifactRepository.get_underlying_uri(model_uri)

File /local_disk0/.ephemeral_nfs/envs/pythonEnv-3ef307dd-1827-4ce9-b993-14ddd54d1d07/lib/python3.9/site-packages/mlflow/utils/file_utils.py:383, in TempDir.__enter__(self)
    382 def __enter__(self):
--> 383     self._path = os.path.abspath(create_tmp_dir())
    384     assert os.path.exists(self._path)
    385     if self._chdr:

File /local_disk0/.ephemeral_nfs/envs/pythonEnv-3ef307dd-1827-4ce9-b993-14ddd54d1d07/lib/python3.9/site-packages/mlflow/utils/file_utils.py:830, in create_tmp_dir()
    827     except Exception:
    828         repl_local_tmp_dir = os.path.join("/tmp", "repl_tmp_data", get_repl_id())
--> 830     return tempfile.mkdtemp(dir=repl_local_tmp_dir)
    831 else:
    832     return tempfile.mkdtemp()

File /usr/lib/python3.9/tempfile.py:498, in mkdtemp(suffix, prefix, dir)
    496 _sys.audit("tempfile.mkdtemp", file)
    497 try:
--> 498     _os.mkdir(file, 0o700)
    499 except FileExistsError:
    500     continue    # try again

FileNotFoundError: [Errno 2] No such file or directory: '/tmp/repl_tmp_data/ReplId-251b4-df166-b640a-3/tmp4sx6kt5d'
```


### Other info / logs

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
REPLACE_ME
```


### What component(s) does this bug affect?

- [ ] `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/gateway`: AI Gateway service, Gateway client APIs, third-party Gateway integrations
- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [X] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/recipes`: Recipes, Recipe APIs, Recipe configs, Recipe Templates
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: MLflow Model server, model deployment tools, Spark UDFs
- [ ] `area/server-infra`: MLflow Tracking server backend
- [ ] `area/tracking`: Tracking Service, tracking client APIs, autologging

### What interface(s) does this bug affect?

- [ ] `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

### What language(s) does this bug affect?

- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

### What integration(s) does this bug affect?

- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations


Discussion:
@tahaum Thanks for the report. Does this error reproduce with MLflow 2.6.0?

Resolving patch:
diff --git a/mlflow/utils/file_utils.py b/mlflow/utils/file_utils.py
--- a/mlflow/utils/file_utils.py
+++ b/mlflow/utils/file_utils.py
@@ -827,6 +827,7 @@ def create_tmp_dir():
         except Exception:
             repl_local_tmp_dir = os.path.join("/tmp", "repl_tmp_data", get_repl_id())
 
+        os.makedirs(repl_local_tmp_dir, exist_ok=True)
         return tempfile.mkdtemp(dir=repl_local_tmp_dir)
     else:
         return tempfile.mkdtemp()


Test patch: