HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

New MLFLOW_KEEP_RUN_ACTIVE flag for enhanced MLflow run management (#8808)

Browse source 
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Massimiliano Riva 2024-03-10 18:28:57 +01:00 • committed by GitHub
parent af6c02c39b
commit 1f9667fff2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 36 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

28
tests/test_integrations.py
View file

@ -29,6 +29,34 @@ def test_mlflow():
SETTINGS["mlflow"] = True
YOLO("yolov8n-cls.yaml").train(data="imagenet10", imgsz=32, epochs=3, plots=False, device="cpu")
@pytest.mark.skipif(not check_requirements('mlflow', install=False), reason='mlflow not installed')
def test_mlflow_keep_run_active():
import os
import mlflow
"""Test training with MLflow tracking enabled."""
SETTINGS['mlflow'] = True
run_name = 'Test Run'
os.environ['MLFLOW_RUN'] = run_name
# Test with MLFLOW_KEEP_RUN_ACTIVE=True
os.environ['MLFLOW_KEEP_RUN_ACTIVE'] = 'True'
YOLO('yolov8n-cls.yaml').train(data='imagenet10', imgsz=32, epochs=1, plots=False, device='cpu')
status = mlflow.active_run().info.status
assert status == 'RUNNING', "MLflow run should be active when MLFLOW_KEEP_RUN_ACTIVE=True"
run_id = mlflow.active_run().info.run_id
# Test with MLFLOW_KEEP_RUN_ACTIVE=False
os.environ['MLFLOW_KEEP_RUN_ACTIVE'] = 'False'
YOLO('yolov8n-cls.yaml').train(data='imagenet10', imgsz=32, epochs=1, plots=False, device='cpu')
status = mlflow.get_run(run_id=run_id).info.status
assert status == 'FINISHED', "MLflow run should be ended when MLFLOW_KEEP_RUN_ACTIVE=False"
# Test with MLFLOW_KEEP_RUN_ACTIVE not set
os.environ.pop('MLFLOW_KEEP_RUN_ACTIVE', None)
YOLO('yolov8n-cls.yaml').train(data='imagenet10', imgsz=32, epochs=1, plots=False, device='cpu')
status = mlflow.get_run(run_id=run_id).info.status
assert status == 'FINISHED', "MLflow run should be ended by default when MLFLOW_KEEP_RUN_ACTIVE is not set"
@pytest.mark.skipif(not check_requirements("tritonclient", install=False), reason="tritonclient[all] not installed")
def test_triton():