1  """System metrics logging module."""
 2  
 3  from mlflow.environment_variables import (
 4      MLFLOW_ENABLE_SYSTEM_METRICS_LOGGING,
 5      MLFLOW_SYSTEM_METRICS_NODE_ID,
 6      MLFLOW_SYSTEM_METRICS_SAMPLES_BEFORE_LOGGING,
 7      MLFLOW_SYSTEM_METRICS_SAMPLING_INTERVAL,
 8  )
 9  
10  
11  def disable_system_metrics_logging():
12      """Disable system metrics logging globally.
13  
14      Calling this function will disable system metrics logging globally, but users can still opt in
15      system metrics logging for individual runs by `mlflow.start_run(log_system_metrics=True)`.
16      """
17      MLFLOW_ENABLE_SYSTEM_METRICS_LOGGING.set(False)
18  
19  
20  def enable_system_metrics_logging():
21      """Enable system metrics logging globally.
22  
23      Calling this function will enable system metrics logging globally, but users can still opt out
24      system metrics logging for individual runs by `mlflow.start_run(log_system_metrics=False)`.
25      """
26      MLFLOW_ENABLE_SYSTEM_METRICS_LOGGING.set(True)
27  
28  
29  def set_system_metrics_sampling_interval(interval):
30      """Set the system metrics sampling interval.
31  
32      Every `interval` seconds, the system metrics will be collected. By default `interval=10`.
33      """
34      if interval is None:
35          MLFLOW_SYSTEM_METRICS_SAMPLING_INTERVAL.unset()
36      else:
37          MLFLOW_SYSTEM_METRICS_SAMPLING_INTERVAL.set(interval)
38  
39  
40  def set_system_metrics_samples_before_logging(samples):
41      """Set the number of samples before logging system metrics.
42  
43      Every time `samples` samples have been collected, the system metrics will be logged to mlflow.
44      By default `samples=1`.
45      """
46      if samples is None:
47          MLFLOW_SYSTEM_METRICS_SAMPLES_BEFORE_LOGGING.unset()
48      else:
49          MLFLOW_SYSTEM_METRICS_SAMPLES_BEFORE_LOGGING.set(samples)
50  
51  
52  def set_system_metrics_node_id(node_id):
53      """Set the system metrics node id.
54  
55      node_id is the identifier of the machine where the metrics are collected. This is useful in
56      multi-node (distributed training) setup.
57      """
58      if node_id is None:
59          MLFLOW_SYSTEM_METRICS_NODE_ID.unset()
60      else:
61          MLFLOW_SYSTEM_METRICS_NODE_ID.set(node_id)