mirror of
https://github.com/THU-MIG/yolov10.git
synced 2025-05-22 21:04:21 +08:00
Fix test_mlflow_keep_run_active()
(#8868)
This commit is contained in:
parent
d24ead447e
commit
3623d62ea8
@ -29,34 +29,38 @@ 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')
|
||||
|
||||
@pytest.mark.skipif(True, reason="Test failing in scheduled CI https://github.com/ultralytics/ultralytics/pull/8868")
|
||||
@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
|
||||
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')
|
||||
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"
|
||||
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')
|
||||
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"
|
||||
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')
|
||||
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"
|
||||
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():
|
||||
|
Loading…
x
Reference in New Issue
Block a user