0
# Configuration and System Management
1
2
MLflow's configuration system provides comprehensive control over tracking URIs, authentication, storage backends, and system behavior. The configuration API enables programmatic setup of MLflow environments with support for various storage systems, authentication methods, and deployment configurations.
3
4
## Capabilities
5
6
### URI and Connection Configuration
7
8
Functions for configuring MLflow tracking and registry URIs with support for various backend storage systems.
9
10
```python { .api }
11
def set_tracking_uri(uri):
12
"""
13
Set the tracking server URI for MLflow operations.
14
15
Parameters:
16
- uri: str - Tracking server URI (file://, http://, https://, databricks://, etc.)
17
18
Supported URI formats:
19
- file:///absolute/path/to/directory (local filesystem)
20
- http://hostname:port (MLflow tracking server)
21
- https://hostname:port (secure MLflow tracking server)
22
- databricks:// (Databricks workspace)
23
- sqlite:///path/to/database.db (SQLite database)
24
- mysql://user:password@host:port/database (MySQL)
25
- postgresql://user:password@host:port/database (PostgreSQL)
26
"""
27
28
def get_tracking_uri():
29
"""
30
Get current MLflow tracking URI.
31
32
Returns:
33
str - Current tracking URI or default if not set
34
"""
35
36
def set_registry_uri(uri):
37
"""
38
Set the model registry URI for MLflow model registry operations.
39
40
Parameters:
41
- uri: str - Model registry URI (same formats as tracking URI)
42
"""
43
44
def get_registry_uri():
45
"""
46
Get current MLflow model registry URI.
47
48
Returns:
49
str - Current registry URI or tracking URI if not set separately
50
"""
51
52
def get_artifact_uri(run_id=None):
53
"""
54
Get artifact storage URI for run or current run.
55
56
Parameters:
57
- run_id: str, optional - Run ID (defaults to current active run)
58
59
Returns:
60
str - Artifact storage URI for the specified run
61
"""
62
63
def set_system_metrics_node_id(node_id):
64
"""
65
Set node identifier for system metrics collection.
66
67
Parameters:
68
- node_id: str - Unique identifier for the current node
69
"""
70
71
def set_system_metrics_sampling_interval(interval):
72
"""
73
Set sampling interval for system metrics collection.
74
75
Parameters:
76
- interval: float - Sampling interval in seconds
77
"""
78
79
def set_system_metrics_samples_before_logging(samples):
80
"""
81
Set number of samples to collect before logging system metrics.
82
83
Parameters:
84
- samples: int - Number of samples to buffer before logging
85
"""
86
```
87
88
### Environment and Authentication
89
90
Functions for managing authentication credentials and environment configuration.
91
92
```python { .api }
93
def set_experiment_id(experiment_id):
94
"""
95
Set active experiment ID for current session.
96
97
Parameters:
98
- experiment_id: str - Experiment ID to set as active
99
"""
100
101
def get_experiment_id():
102
"""
103
Get current active experiment ID.
104
105
Returns:
106
str - Active experiment ID or None if not set
107
"""
108
109
def set_run_id(run_id):
110
"""
111
Set active run ID for current session.
112
113
Parameters:
114
- run_id: str - Run ID to set as active
115
"""
116
117
def get_run_id():
118
"""
119
Get current active run ID.
120
121
Returns:
122
str - Active run ID or None if not set
123
"""
124
```
125
126
### Artifact Storage Configuration
127
128
Configuration options for artifact storage backends including S3, Azure, GCS, and local storage.
129
130
```python { .api }
131
def set_default_artifact_root(artifact_root):
132
"""
133
Set default artifact root for experiments.
134
135
Parameters:
136
- artifact_root: str - Default artifact storage location
137
138
Supported formats:
139
- s3://bucket/path (Amazon S3)
140
- gs://bucket/path (Google Cloud Storage)
141
- wasbs://container@account.blob.core.windows.net/path (Azure Blob)
142
- hdfs://namenode:port/path (HDFS)
143
- file:///absolute/path (local filesystem)
144
"""
145
146
def get_default_artifact_root():
147
"""
148
Get default artifact root configuration.
149
150
Returns:
151
str - Default artifact root path
152
"""
153
```
154
155
### System Metrics Configuration
156
157
Configuration for automatic system metrics collection including CPU, memory, and GPU usage.
158
159
```python { .api }
160
def enable_system_metrics_logging():
161
"""
162
Enable automatic system metrics logging for MLflow runs.
163
Collects CPU, memory, disk, and network metrics during run execution.
164
"""
165
166
def disable_system_metrics_logging():
167
"""
168
Disable automatic system metrics logging.
169
"""
170
171
def is_system_metrics_logging_enabled():
172
"""
173
Check if system metrics logging is enabled.
174
175
Returns:
176
bool - True if system metrics logging is enabled
177
"""
178
```
179
180
### Backend Store Configuration
181
182
Configuration utilities for different backend store types and database connections.
183
184
```python { .api }
185
def is_tracking_uri_set():
186
"""
187
Check if tracking URI has been explicitly set.
188
189
Returns:
190
bool - True if tracking URI is explicitly configured
191
"""
192
193
def get_db_profile_from_uri(db_uri):
194
"""
195
Extract database connection profile from URI.
196
197
Parameters:
198
- db_uri: str - Database URI
199
200
Returns:
201
dict - Database connection profile with host, port, database, etc.
202
"""
203
204
def construct_db_uri_from_profile(profile):
205
"""
206
Construct database URI from connection profile.
207
208
Parameters:
209
- profile: dict - Database connection profile
210
211
Returns:
212
str - Constructed database URI
213
"""
214
```
215
216
### Configuration Validation
217
218
Functions for validating configuration settings and checking system requirements.
219
220
```python { .api }
221
def validate_tracking_uri(uri):
222
"""
223
Validate tracking URI format and accessibility.
224
225
Parameters:
226
- uri: str - Tracking URI to validate
227
228
Returns:
229
bool - True if URI is valid and accessible
230
231
Raises:
232
MlflowException - If URI is invalid or inaccessible
233
"""
234
235
def validate_artifact_root(artifact_root):
236
"""
237
Validate artifact storage root configuration.
238
239
Parameters:
240
- artifact_root: str - Artifact root path to validate
241
242
Returns:
243
bool - True if artifact root is valid and accessible
244
"""
245
246
def check_server_connection(tracking_uri=None):
247
"""
248
Check connection to MLflow tracking server.
249
250
Parameters:
251
- tracking_uri: str, optional - Server URI to check (defaults to current)
252
253
Returns:
254
dict - Server status and version information
255
"""
256
```
257
258
### Environment Variables and Settings
259
260
Configuration through environment variables and programmatic settings management.
261
262
```python { .api }
263
class MLflowEnvironment:
264
"""
265
MLflow environment variable constants and utilities.
266
"""
267
268
# Tracking and Registry URIs
269
MLFLOW_TRACKING_URI = "MLFLOW_TRACKING_URI"
270
MLFLOW_REGISTRY_URI = "MLFLOW_REGISTRY_URI"
271
272
# Authentication
273
MLFLOW_TRACKING_USERNAME = "MLFLOW_TRACKING_USERNAME"
274
MLFLOW_TRACKING_PASSWORD = "MLFLOW_TRACKING_PASSWORD"
275
MLFLOW_TRACKING_TOKEN = "MLFLOW_TRACKING_TOKEN"
276
MLFLOW_TRACKING_INSECURE_TLS = "MLFLOW_TRACKING_INSECURE_TLS"
277
278
# Artifact Storage
279
MLFLOW_DEFAULT_ARTIFACT_ROOT = "MLFLOW_DEFAULT_ARTIFACT_ROOT"
280
MLFLOW_ARTIFACT_UPLOAD_DOWNLOAD_TIMEOUT = "MLFLOW_ARTIFACT_UPLOAD_DOWNLOAD_TIMEOUT"
281
282
# System Configuration
283
MLFLOW_EXPERIMENT_ID = "MLFLOW_EXPERIMENT_ID"
284
MLFLOW_RUN_ID = "MLFLOW_RUN_ID"
285
MLFLOW_ENABLE_SYSTEM_METRICS_LOGGING = "MLFLOW_ENABLE_SYSTEM_METRICS_LOGGING"
286
287
# S3 Configuration
288
AWS_ACCESS_KEY_ID = "AWS_ACCESS_KEY_ID"
289
AWS_SECRET_ACCESS_KEY = "AWS_SECRET_ACCESS_KEY"
290
AWS_SESSION_TOKEN = "AWS_SESSION_TOKEN"
291
MLFLOW_S3_ENDPOINT_URL = "MLFLOW_S3_ENDPOINT_URL"
292
MLFLOW_S3_IGNORE_TLS = "MLFLOW_S3_IGNORE_TLS"
293
294
# Azure Configuration
295
AZURE_STORAGE_CONNECTION_STRING = "AZURE_STORAGE_CONNECTION_STRING"
296
AZURE_STORAGE_ACCESS_KEY = "AZURE_STORAGE_ACCESS_KEY"
297
298
# Google Cloud Configuration
299
GOOGLE_APPLICATION_CREDENTIALS = "GOOGLE_APPLICATION_CREDENTIALS"
300
301
@staticmethod
302
def get_env_var(key, default=None):
303
"""Get environment variable value."""
304
import os
305
return os.getenv(key, default)
306
307
@staticmethod
308
def set_env_var(key, value):
309
"""Set environment variable value."""
310
import os
311
os.environ[key] = str(value)
312
313
def get_env_vars_dict():
314
"""
315
Get dictionary of all MLflow-related environment variables.
316
317
Returns:
318
dict - Environment variables and their values
319
"""
320
321
def set_env_vars_from_dict(env_dict):
322
"""
323
Set environment variables from dictionary.
324
325
Parameters:
326
- env_dict: dict - Environment variables to set
327
"""
328
```
329
330
### Configuration Context Managers
331
332
Context managers for temporary configuration changes and isolated environments.
333
334
```python { .api }
335
@contextmanager
336
def temporary_tracking_uri(uri):
337
"""
338
Temporarily set tracking URI for block execution.
339
340
Parameters:
341
- uri: str - Temporary tracking URI
342
343
Usage:
344
with temporary_tracking_uri("sqlite:///temp.db"):
345
# Operations use temporary URI
346
mlflow.create_experiment("temp_exp")
347
# Original URI restored
348
"""
349
350
@contextmanager
351
def temporary_experiment_id(experiment_id):
352
"""
353
Temporarily set experiment ID for block execution.
354
355
Parameters:
356
- experiment_id: str - Temporary experiment ID
357
"""
358
359
@contextmanager
360
def isolated_mlflow_config():
361
"""
362
Create isolated MLflow configuration context.
363
All configuration changes are isolated and restored on exit.
364
"""
365
```
366
367
## Usage Examples
368
369
### Basic Configuration Setup
370
371
```python
372
import mlflow
373
import os
374
375
# Set tracking URI
376
mlflow.set_tracking_uri("http://localhost:5000")
377
378
# Verify connection
379
print(f"Tracking URI: {mlflow.get_tracking_uri()}")
380
381
# Set registry URI (if different from tracking)
382
mlflow.set_registry_uri("sqlite:///model_registry.db")
383
384
# Set default artifact location
385
mlflow.set_default_artifact_root("s3://my-mlflow-bucket/artifacts")
386
387
# Enable system metrics logging
388
mlflow.enable_system_metrics_logging()
389
390
# Check current configuration
391
print(f"Registry URI: {mlflow.get_registry_uri()}")
392
print(f"Artifact root: {mlflow.get_default_artifact_root()}")
393
print(f"System metrics enabled: {mlflow.is_system_metrics_logging_enabled()}")
394
```
395
396
### Environment Variable Configuration
397
398
```python
399
import os
400
import mlflow
401
402
# Set configuration via environment variables
403
os.environ["MLFLOW_TRACKING_URI"] = "https://mlflow.company.com"
404
os.environ["MLFLOW_TRACKING_USERNAME"] = "user@company.com"
405
os.environ["MLFLOW_TRACKING_PASSWORD"] = "secure_password"
406
os.environ["MLFLOW_DEFAULT_ARTIFACT_ROOT"] = "s3://company-mlflow/artifacts"
407
os.environ["MLFLOW_ENABLE_SYSTEM_METRICS_LOGGING"] = "true"
408
409
# MLflow automatically uses environment variables
410
print(f"Tracking URI: {mlflow.get_tracking_uri()}")
411
412
# Alternative: Use MLflowEnvironment helper
413
from mlflow.config import MLflowEnvironment
414
415
MLflowEnvironment.set_env_var("MLFLOW_EXPERIMENT_ID", "123")
416
experiment_id = MLflowEnvironment.get_env_var("MLFLOW_EXPERIMENT_ID")
417
print(f"Experiment ID: {experiment_id}")
418
419
# Get all MLflow environment variables
420
env_vars = mlflow.get_env_vars_dict()
421
for key, value in env_vars.items():
422
if value: # Only show set variables
423
print(f"{key}: {value}")
424
```
425
426
### Database Backend Configuration
427
428
```python
429
import mlflow
430
431
# PostgreSQL backend
432
postgres_uri = "postgresql://mlflow:password@localhost:5432/mlflow_db"
433
mlflow.set_tracking_uri(postgres_uri)
434
435
# MySQL backend
436
mysql_uri = "mysql://mlflow:password@localhost:3306/mlflow_db"
437
mlflow.set_tracking_uri(mysql_uri)
438
439
# SQLite backend (for development)
440
sqlite_uri = "sqlite:///mlflow.db"
441
mlflow.set_tracking_uri(sqlite_uri)
442
443
# Validate configuration
444
try:
445
mlflow.validate_tracking_uri(mlflow.get_tracking_uri())
446
print("Tracking URI is valid")
447
448
# Check server connection
449
server_status = mlflow.check_server_connection()
450
print(f"Server version: {server_status.get('version')}")
451
print(f"Server status: {server_status.get('status')}")
452
453
except Exception as e:
454
print(f"Configuration error: {e}")
455
```
456
457
### Cloud Storage Configuration
458
459
```python
460
import mlflow
461
import os
462
463
# AWS S3 Configuration
464
os.environ["AWS_ACCESS_KEY_ID"] = "your_access_key"
465
os.environ["AWS_SECRET_ACCESS_KEY"] = "your_secret_key"
466
os.environ["MLFLOW_S3_ENDPOINT_URL"] = "https://s3.amazonaws.com"
467
468
mlflow.set_tracking_uri("http://localhost:5000")
469
mlflow.set_default_artifact_root("s3://my-mlflow-bucket/artifacts")
470
471
# Azure Blob Storage Configuration
472
os.environ["AZURE_STORAGE_CONNECTION_STRING"] = "DefaultEndpointsProtocol=https;..."
473
mlflow.set_default_artifact_root("wasbs://container@account.blob.core.windows.net/artifacts")
474
475
# Google Cloud Storage Configuration
476
os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = "/path/to/service-account.json"
477
mlflow.set_default_artifact_root("gs://my-mlflow-bucket/artifacts")
478
479
# Validate artifact storage
480
try:
481
mlflow.validate_artifact_root(mlflow.get_default_artifact_root())
482
print("Artifact root is accessible")
483
except Exception as e:
484
print(f"Artifact storage error: {e}")
485
```
486
487
### System Metrics Configuration
488
489
```python
490
import mlflow
491
492
# Configure system metrics collection
493
mlflow.set_system_metrics_node_id("worker-node-1")
494
mlflow.set_system_metrics_sampling_interval(1.0) # 1 second intervals
495
mlflow.set_system_metrics_samples_before_logging(60) # Log every 60 samples
496
497
# Enable system metrics
498
mlflow.enable_system_metrics_logging()
499
500
# Start experiment with system metrics
501
mlflow.set_experiment("system-metrics-experiment")
502
503
with mlflow.start_run():
504
# Run your ML code - system metrics collected automatically
505
import time
506
import numpy as np
507
508
# Simulate ML training workload
509
for epoch in range(10):
510
# Simulate CPU/memory intensive work
511
data = np.random.randn(10000, 100)
512
result = np.dot(data, data.T)
513
514
mlflow.log_metric("epoch", epoch)
515
time.sleep(2) # System metrics collected during sleep
516
517
print("System metrics logged automatically")
518
519
# Check system metrics in MLflow UI
520
```
521
522
### Temporary Configuration Context
523
524
```python
525
import mlflow
526
from mlflow.config import temporary_tracking_uri, temporary_experiment_id
527
528
# Set default configuration
529
mlflow.set_tracking_uri("http://localhost:5000")
530
mlflow.set_experiment("default-experiment")
531
532
# Use temporary configuration
533
with temporary_tracking_uri("sqlite:///temp_experiments.db"):
534
# This block uses temporary SQLite backend
535
temp_exp_id = mlflow.create_experiment("temporary-experiment")
536
537
with mlflow.start_run():
538
mlflow.log_metric("temp_metric", 42)
539
print("Logged to temporary backend")
540
541
# Back to original configuration
542
print(f"Original tracking URI: {mlflow.get_tracking_uri()}")
543
544
# Temporary experiment ID
545
with temporary_experiment_id("123"):
546
with mlflow.start_run():
547
mlflow.log_metric("experiment_123_metric", 100)
548
print("Logged to experiment 123")
549
550
# Original experiment restored
551
print(f"Current experiment: {mlflow.get_experiment_id()}")
552
```
553
554
### Configuration Validation and Debugging
555
556
```python
557
import mlflow
558
from mlflow.config import MLflowEnvironment
559
560
def diagnose_mlflow_config():
561
"""Comprehensive MLflow configuration diagnostic."""
562
563
print("=== MLflow Configuration Diagnostic ===")
564
565
# Check basic configuration
566
print(f"Tracking URI: {mlflow.get_tracking_uri()}")
567
print(f"Registry URI: {mlflow.get_registry_uri()}")
568
print(f"Artifact Root: {mlflow.get_default_artifact_root()}")
569
570
# Check environment variables
571
print("\n=== Environment Variables ===")
572
env_vars = mlflow.get_env_vars_dict()
573
for key, value in sorted(env_vars.items()):
574
if value:
575
# Mask sensitive values
576
if "password" in key.lower() or "token" in key.lower() or "key" in key.lower():
577
value = "***masked***"
578
print(f"{key}: {value}")
579
580
# Validate configuration
581
print("\n=== Validation ===")
582
try:
583
mlflow.validate_tracking_uri(mlflow.get_tracking_uri())
584
print("✓ Tracking URI is valid")
585
except Exception as e:
586
print(f"✗ Tracking URI error: {e}")
587
588
try:
589
mlflow.validate_artifact_root(mlflow.get_default_artifact_root())
590
print("✓ Artifact root is accessible")
591
except Exception as e:
592
print(f"✗ Artifact root error: {e}")
593
594
# Test server connection
595
try:
596
server_info = mlflow.check_server_connection()
597
print(f"✓ Server connection successful")
598
print(f" Version: {server_info.get('version', 'unknown')}")
599
except Exception as e:
600
print(f"✗ Server connection failed: {e}")
601
602
# System metrics status
603
print(f"\nSystem metrics enabled: {mlflow.is_system_metrics_logging_enabled()}")
604
605
# Run diagnostic
606
diagnose_mlflow_config()
607
```
608
609
### Advanced Configuration Management
610
611
```python
612
import mlflow
613
from mlflow.config import MLflowEnvironment, isolated_mlflow_config
614
import json
615
616
class MLflowConfigManager:
617
"""Advanced MLflow configuration management."""
618
619
def __init__(self):
620
self.configs = {}
621
622
def save_config(self, name, config_dict):
623
"""Save configuration profile."""
624
self.configs[name] = config_dict
625
626
def load_config(self, name):
627
"""Load and apply configuration profile."""
628
if name not in self.configs:
629
raise ValueError(f"Configuration '{name}' not found")
630
631
config = self.configs[name]
632
633
if "tracking_uri" in config:
634
mlflow.set_tracking_uri(config["tracking_uri"])
635
if "registry_uri" in config:
636
mlflow.set_registry_uri(config["registry_uri"])
637
if "artifact_root" in config:
638
mlflow.set_default_artifact_root(config["artifact_root"])
639
if "system_metrics" in config:
640
if config["system_metrics"]:
641
mlflow.enable_system_metrics_logging()
642
else:
643
mlflow.disable_system_metrics_logging()
644
645
def export_config(self, name, filepath):
646
"""Export configuration to file."""
647
if name in self.configs:
648
with open(filepath, 'w') as f:
649
json.dump(self.configs[name], f, indent=2)
650
651
def import_config(self, name, filepath):
652
"""Import configuration from file."""
653
with open(filepath, 'r') as f:
654
self.configs[name] = json.load(f)
655
656
# Usage example
657
config_manager = MLflowConfigManager()
658
659
# Define configurations for different environments
660
config_manager.save_config("development", {
661
"tracking_uri": "sqlite:///dev_mlflow.db",
662
"artifact_root": "./dev_artifacts",
663
"system_metrics": True
664
})
665
666
config_manager.save_config("staging", {
667
"tracking_uri": "http://staging-mlflow:5000",
668
"registry_uri": "postgresql://user:pass@staging-db:5432/mlflow",
669
"artifact_root": "s3://staging-mlflow-bucket/artifacts",
670
"system_metrics": False
671
})
672
673
config_manager.save_config("production", {
674
"tracking_uri": "https://prod-mlflow.company.com",
675
"registry_uri": "postgresql://user:pass@prod-db:5432/mlflow",
676
"artifact_root": "s3://prod-mlflow-bucket/artifacts",
677
"system_metrics": True
678
})
679
680
# Switch between configurations
681
config_manager.load_config("development")
682
print(f"Development config loaded: {mlflow.get_tracking_uri()}")
683
684
config_manager.load_config("production")
685
print(f"Production config loaded: {mlflow.get_tracking_uri()}")
686
687
# Export configuration for sharing
688
config_manager.export_config("production", "prod_config.json")
689
```
690
691
## Types
692
693
```python { .api }
694
from typing import Dict, Any, Optional, ContextManager
695
import os
696
697
class MLflowEnvironment:
698
# Environment variable constants
699
MLFLOW_TRACKING_URI: str
700
MLFLOW_REGISTRY_URI: str
701
MLFLOW_TRACKING_USERNAME: str
702
MLFLOW_TRACKING_PASSWORD: str
703
MLFLOW_TRACKING_TOKEN: str
704
MLFLOW_DEFAULT_ARTIFACT_ROOT: str
705
MLFLOW_EXPERIMENT_ID: str
706
MLFLOW_RUN_ID: str
707
MLFLOW_ENABLE_SYSTEM_METRICS_LOGGING: str
708
709
@staticmethod
710
def get_env_var(key: str, default: Optional[str] = None) -> Optional[str]: ...
711
712
@staticmethod
713
def set_env_var(key: str, value: str) -> None: ...
714
715
# Configuration functions return types
716
def get_tracking_uri() -> str: ...
717
def get_registry_uri() -> str: ...
718
def get_artifact_uri(run_id: Optional[str] = None) -> str: ...
719
def get_default_artifact_root() -> str: ...
720
def get_experiment_id() -> Optional[str]: ...
721
def get_run_id() -> Optional[str]: ...
722
723
def is_tracking_uri_set() -> bool: ...
724
def is_system_metrics_logging_enabled() -> bool: ...
725
726
def validate_tracking_uri(uri: str) -> bool: ...
727
def validate_artifact_root(artifact_root: str) -> bool: ...
728
729
def check_server_connection(tracking_uri: Optional[str] = None) -> Dict[str, Any]: ...
730
def get_env_vars_dict() -> Dict[str, str]: ...
731
732
# Context managers
733
def temporary_tracking_uri(uri: str) -> ContextManager[None]: ...
734
def temporary_experiment_id(experiment_id: str) -> ContextManager[None]: ...
735
def isolated_mlflow_config() -> ContextManager[None]: ...
736
737
# Configuration profile type
738
ConfigProfile = Dict[str, Any]
739
740
# Database profile type
741
DatabaseProfile = Dict[str, str] # host, port, database, username, etc.
742
```