Repository: mlflow/mlflow

Issue:
[BUG] "log_system_metrics" not working when "existing_run_id" is set
### 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)

### Where did you encounter this bug?

Local machine

### Willingness to contribute

Yes. I can contribute a fix for this bug independently.

### MLflow version

- Client: 2.8.0
- Tracking server: 2.8.0

### System information

- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: Linux Ubuntu 22.04
- **Python version**: 3.10.12


### Describe the problem

### Issue
When calling `mlflow.start_run` with `log_system_metrics=True` and `run_id` set by parameter or environment variable, the `SystemMetricsMonitor` is not going to be started. System metrics are only recorded for newly created runs, while they are not for continuing existing runs. This is a bug, since one would expect that system metrics are recorded in both cases.

This bug might be unintended due to an incorrect indentation of the code block in `mlflow/tracking/fluent.py` which starts the monitoring process only in the else-branch (i.e. when `existing_run_id` is not set).
![image](https://github.com/mlflow/mlflow/assets/44120873/aa7f5937-462c-4044-9ecf-bc964871597c)

It might also be intended in order to avoid dealing with the question from which `step` the recording of the metrics should continue for existing runs. When moving the code block outside of the else-branch, system metrics are going to be logged for existing runs with the current timestamp, but starting from step 0 again. Hence, the plot with the time on the x-axis is still correct, but not the plot with the steps on the x-axis.
![image](https://github.com/mlflow/mlflow/assets/44120873/779cc7c4-ee40-4169-bb47-e9a8c3d90716)
![image](https://github.com/mlflow/mlflow/assets/44120873/90067f56-50c8-4020-a980-4e284930ea67)

There is also a minor bug in `system_metrics_monitor.py`, which I have found when inspecting the other issue. Since the system metrics are logged with `BatchMetricsLogger`, it is advisable to call `self.mlflow_logger.flush` when stopping the monitoring process.

### Proposed Solution
My suggestion is to move the code block (see above) out of the else-branch in any case, as it still provides reasonable plots with the time on the x-axis and having the issue with the steps on the x-axis is acceptable for an experimental feature in my opinion.

For solving the issue with the steps on the x-axis, I can think of different solutions for setting `_logging_step` in the constructor of `SystemMetricsLogger`:
- determine the largest step for one of the system metrics from which the logging should be continued using  `mlflow.client.MlflowClient().get_metric_history(<run_id>, <metric_key>)` (The question is which metric should be taken? In particular, it has to be one which is always logged, i.e. not `system/gpu...`.)
- determine the largest step among all system metrics from which the logging should be continued (This is more robust, in case some system metrics are not logged and some are. However, it might be inefficient to iterate through all metric histories...)
- log an additional metric `system/_logging_step` from which the logging should be continued (Not my favorite one...)

One can also think of extending this to have better looking plots:
- add an offset of X steps when continuing a run in order to artificially create some distance between the two runs in the plots
- log a NaN value first when continuing a runs in order to artificially create a gap between the two runs in the plot

But here I need more input from the maintainers/community regarding the design of a fix.

Regarding the minor issue, I would simply add `self.mlflow_logger.flush()` to the end of `SystemMetricsMonitor.finish` once all monitoring processes have been stopped.

PS: That's my first Github activity, so I hope I did everything correct:D Let me know if I should provide more details/do something differently.

### Tracking information

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```shell
System information: Linux #1 SMP Fri Jan 27 02:56:13 UTC 2023
Python version: 3.10.12
MLflow version: 2.8.0
MLflow module location: x
Tracking URI: x
Registry URI: x
MLflow dependencies: 
  Flask: 2.3.3
  Jinja2: 3.1.2
  alembic: 1.12.0
  click: 8.1.7
  cloudpickle: 2.2.1
  databricks-cli: 0.18.0
  docker: 6.1.3
  entrypoints: 0.4
  gitpython: 3.1.40
  gunicorn: 21.2.0
  importlib-metadata: 6.8.0
  markdown: 3.5
  matplotlib: 3.8.0
  numpy: 1.25.2
  packaging: 23.2
  pandas: 2.1.1
  protobuf: 4.24.4
  psutil: 5.9.6
  pyarrow: 13.0.0
  pytz: 2023.3.post1
  pyyaml: 6.0.1
  querystring-parser: 1.2.4
  requests: 2.31.0
  scikit-learn: 1.3.1
  scipy: 1.11.3
  sqlalchemy: 2.0.22
  sqlparse: 0.4.4
```


### Code to reproduce issue

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
import mlflow
import time

mlflow.set_system_metrics_sampling_interval(1)

with mlflow.start_run(log_system_metrics=True) as run:
    time.sleep(2)

with mlflow.start_run(log_system_metrics=True, run_id=run.info.run_id):
    time.sleep(2)
```


### Stack trace

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
*This bug does not generate an exception, but it will become visible in the log (see below).*

### Other info / logs

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
Expected output:
```
2023/10/31 18:53:45 INFO mlflow.system_metrics.system_metrics_monitor: Started monitoring system metrics.
2023/10/31 18:53:47 INFO mlflow.system_metrics.system_metrics_monitor: Stopping system metrics monitoring...
2023/10/31 18:53:47 INFO mlflow.system_metrics.system_metrics_monitor: Successfully terminated system metrics monitoring!
2023/10/31 18:53:47 INFO mlflow.system_metrics.system_metrics_monitor: Started monitoring system metrics.
2023/10/31 18:53:49 INFO mlflow.system_metrics.system_metrics_monitor: Stopping system metrics monitoring...
2023/10/31 18:53:49 INFO mlflow.system_metrics.system_metrics_monitor: Successfully terminated system metrics monitoring!
```

Actual output:
```
2023/10/31 18:53:45 INFO mlflow.system_metrics.system_metrics_monitor: Started monitoring system metrics.
2023/10/31 18:53:47 INFO mlflow.system_metrics.system_metrics_monitor: Stopping system metrics monitoring...
2023/10/31 18:53:47 INFO mlflow.system_metrics.system_metrics_monitor: Successfully terminated system metrics monitoring!
```


### 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
- [ ] `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
- [X] `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:
@chenmoneygithub Can you take a look?
@danielyxyang Thanks for reporting the issue! Also extra kudos to the detailed information, it's very clear and helpful! 

Yes, we should start system metrics logging correctly for existing runs. Would you like to contribute a PR for that?

About resuming the steps, we should be able to use the step information from `cpu_utilization_percentage`, which is always logged. Please let me know if you want to work on it, otherwise I can fix that. Thanks!
Sure, I can do that in the next few days once I find some time for it :)

Resolving patch:
diff --git a/mlflow/system_metrics/system_metrics_monitor.py b/mlflow/system_metrics/system_metrics_monitor.py
--- a/mlflow/system_metrics/system_metrics_monitor.py
+++ b/mlflow/system_metrics/system_metrics_monitor.py
@@ -7,6 +7,7 @@
     MLFLOW_SYSTEM_METRICS_SAMPLES_BEFORE_LOGGING,
     MLFLOW_SYSTEM_METRICS_SAMPLING_INTERVAL,
 )
+from mlflow.exceptions import MlflowException
 from mlflow.system_metrics.metrics.cpu_monitor import CPUMonitor
 from mlflow.system_metrics.metrics.disk_monitor import DiskMonitor
 from mlflow.system_metrics.metrics.gpu_monitor import GPUMonitor
@@ -36,9 +37,18 @@ class SystemMetricsMonitor:
         samples_before_logging: int, default to 1. The number of samples to aggregate before
             logging. Will be overridden by `MLFLOW_SYSTEM_METRICS_SAMPLES_BEFORE_LOGGING`
             evnironment variable.
+        resume_logging: bool, default to False. If True, we will resume the system metrics logging
+            from the `run_id`, and the first step to log will be the last step of `run_id` + 1, if
+            False, system metrics logging will start from step 0.
     """
 
-    def __init__(self, run_id, sampling_interval=10, samples_before_logging=1):
+    def __init__(
+        self,
+        run_id,
+        sampling_interval=10,
+        samples_before_logging=1,
+        resume_logging=False,
+    ):
         from mlflow.utils.autologging_utils import BatchMetricsLogger
 
         # Instantiate default monitors.
@@ -60,8 +70,26 @@ def __init__(self, run_id, sampling_interval=10, samples_before_logging=1):
         self.mlflow_logger = BatchMetricsLogger(self._run_id)
         self._shutdown_event = threading.Event()
         self._process = None
-        self._logging_step = 0
         self._metrics_prefix = "system/"
+        self._logging_step = self._get_next_logging_step(run_id) if resume_logging else 0
+
+    def _get_next_logging_step(self, run_id):
+        from mlflow.tracking.client import MlflowClient
+
+        client = MlflowClient()
+        try:
+            run = client.get_run(run_id)
+        except MlflowException:
+            return 0
+        system_metric_name = None
+        for metric_name in run.data.metrics.keys():
+            if metric_name.startswith(self._metrics_prefix):
+                system_metric_name = metric_name
+                break
+        if system_metric_name is None:
+            return 0
+        metric_history = client.get_metric_history(run_id, system_metric_name)
+        return metric_history[-1].step + 1
 
     def start(self):
         """Start monitoring system metrics."""
diff --git a/mlflow/tracking/fluent.py b/mlflow/tracking/fluent.py
--- a/mlflow/tracking/fluent.py
+++ b/mlflow/tracking/fluent.py
@@ -376,32 +376,32 @@ def start_run(
             tags=resolved_tags,
             run_name=run_name,
         )
-        if log_system_metrics is None:
-            # if `log_system_metrics` is not specified, we will check environment variable.
-            log_system_metrics = MLFLOW_ENABLE_SYSTEM_METRICS_LOGGING.get()
-        if log_system_metrics:
-            # Ensure psutil is installed. It was moved to an optional dependency, as it doesn't
-            # have binary for Arm64 Linux and requires build from CPython which is a headache.
-            # https://github.com/giampaolo/psutil/issues/1972
-            if importlib.util.find_spec("psutil") is None:
-                raise MlflowException(
-                    "Failed to start system metrics monitoring as package `psutil` is not "
-                    "installed. Run `pip install psutil` to resolve the issue, "
-                    "otherwise you can disable system metrics logging by passing "
-                    "`log_system_metrics=False` to `start_run()` or setting environment "
-                    f"variable {MLFLOW_ENABLE_SYSTEM_METRICS_LOGGING} to False."
-                )
 
-            try:
-                from mlflow.system_metrics.system_metrics_monitor import SystemMetricsMonitor
+    if log_system_metrics is None:
+        # If `log_system_metrics` is not specified, we will check environment variable.
+        log_system_metrics = MLFLOW_ENABLE_SYSTEM_METRICS_LOGGING.get()
+
+    if log_system_metrics:
+        if importlib.util.find_spec("psutil") is None:
+            raise MlflowException(
+                "Failed to start system metrics monitoring as package `psutil` is not installed. "
+                "Please run `pip install psutil` to resolve the issue, otherwise you can disable "
+                "system metrics logging by passing `log_system_metrics=False` to "
+                "`mlflow.start_run()` or calling `mlflow.disable_system_metrics_logging`."
+            )
+        try:
+            from mlflow.system_metrics.system_metrics_monitor import SystemMetricsMonitor
 
-                system_monitor = SystemMetricsMonitor(active_run_obj.info.run_id)
-                global run_id_to_system_metrics_monitor
+            system_monitor = SystemMetricsMonitor(
+                active_run_obj.info.run_id,
+                resume_logging=existing_run_id is not None,
+            )
+            global run_id_to_system_metrics_monitor
 
-                run_id_to_system_metrics_monitor[active_run_obj.info.run_id] = system_monitor
-                system_monitor.start()
-            except Exception as e:
-                _logger.error(f"Failed to start system metrics monitoring: {e}.")
+            run_id_to_system_metrics_monitor[active_run_obj.info.run_id] = system_monitor
+            system_monitor.start()
+        except Exception as e:
+            _logger.error(f"Failed to start system metrics monitoring: {e}.")
 
     _active_run_stack.append(ActiveRun(active_run_obj))
     return _active_run_stack[-1]


Test patch:
diff --git a/tests/system_metrics/test_system_metrics_logging.py b/tests/system_metrics/test_system_metrics_logging.py
--- a/tests/system_metrics/test_system_metrics_logging.py
+++ b/tests/system_metrics/test_system_metrics_logging.py
@@ -1,10 +1,21 @@
 import threading
 import time
 
+import pytest
+
 import mlflow
 from mlflow.system_metrics.system_metrics_monitor import SystemMetricsMonitor
 
 
+@pytest.fixture(autouse=True)
+def disable_system_metrics_logging():
+    yield
+    # Unset the environment variables to avoid affecting other test cases.
+    mlflow.disable_system_metrics_logging()
+    mlflow.set_system_metrics_sampling_interval(None)
+    mlflow.set_system_metrics_samples_before_logging(None)
+
+
 def test_manual_system_metrics_monitor():
     with mlflow.start_run(log_system_metrics=False) as run:
         system_monitor = SystemMetricsMonitor(
@@ -41,7 +52,7 @@ def test_manual_system_metrics_monitor():
         assert name in metrics
 
     # Check the step is correctly logged.
-    metrics_history = mlflow.tracking.MlflowClient().get_metric_history(
+    metrics_history = mlflow.MlflowClient().get_metric_history(
         run.info.run_id, "system/cpu_utilization_percentage"
     )
     assert metrics_history[-1].step > 0
@@ -82,12 +93,51 @@ def test_automatic_system_metrics_monitor():
         assert name in metrics
 
     # Check the step is correctly logged.
-    metrics_history = mlflow.tracking.MlflowClient().get_metric_history(
+    metrics_history = mlflow.MlflowClient().get_metric_history(
         run.info.run_id, "system/cpu_utilization_percentage"
     )
     assert metrics_history[-1].step > 0
 
-    # Unset the environment variables to avoid affecting other test cases.
-    mlflow.disable_system_metrics_logging()
-    mlflow.set_system_metrics_sampling_interval(None)
-    mlflow.set_system_metrics_samples_before_logging(None)
+
+def test_automatic_system_metrics_monitor_resume_existing_run():
+    mlflow.enable_system_metrics_logging()
+    mlflow.set_system_metrics_sampling_interval(0.2)
+    mlflow.set_system_metrics_samples_before_logging(2)
+    with mlflow.start_run() as run:
+        time.sleep(2)
+
+    # Pause for a bit to allow the system metrics monitoring to exit.
+    time.sleep(1)
+    thread_names = [thread.name for thread in threading.enumerate()]
+    # Check the system metrics monitoring thread has exited.
+    assert "SystemMetricsMonitor" not in thread_names
+
+    # Get the last step.
+    metrics_history = mlflow.MlflowClient().get_metric_history(
+        run.info.run_id, "system/cpu_utilization_percentage"
+    )
+    last_step = metrics_history[-1].step
+
+    with mlflow.start_run(run.info.run_id) as run:
+        time.sleep(2)
+    mlflow_run = mlflow.get_run(run.info.run_id)
+    metrics = mlflow_run.data.metrics
+
+    expected_metrics_name = [
+        "cpu_utilization_percentage",
+        "system_memory_usage_megabytes",
+        "disk_usage_percentage",
+        "disk_usage_megabytes",
+        "disk_available_megabytes",
+        "network_receive_megabytes",
+        "network_transmit_megabytes",
+    ]
+    expected_metrics_name = [f"system/{name}" for name in expected_metrics_name]
+    for name in expected_metrics_name:
+        assert name in metrics
+
+    # Check the step is correctly resumed.
+    metrics_history = mlflow.MlflowClient().get_metric_history(
+        run.info.run_id, "system/cpu_utilization_percentage"
+    )
+    assert metrics_history[-1].step > last_step