Skip to content

Commit

Permalink
fix(test_tuning): added extra tests.
Browse files Browse the repository at this point in the history
  • Loading branch information
mathysgrapotte committed Feb 4, 2025
1 parent 6f6471e commit 503048a
Showing 1 changed file with 17 additions and 9 deletions.
26 changes: 17 additions & 9 deletions tests/cli/test_tuning.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
"""Test the tuning CLI."""

import os
import shutil
import warnings
from pathlib import Path

Expand Down Expand Up @@ -55,27 +56,34 @@ def test_tuning_main(data_path: str, data_config: str, model_path: str, model_co
assert os.path.exists(model_config), f"Model config not found at {model_config}"

try:
# Run main function - should complete without errors
results_dir = Path("tests/test_data/titanic/test_results/").resolve()
results_dir.mkdir(parents=True, exist_ok=True)

# Use directory path for Ray results and file paths for outputs
tuning.main(
model_path=model_path,
data_path=data_path,
data_config_path=data_config,
model_config_path=model_config,
initial_weights=None,
ray_results_dirpath=None,
output_path=None,
best_optimizer_path=None,
best_metrics_path=None,
best_config_path=None,
ray_results_dirpath=str(results_dir), # Directory path without URI scheme
output_path=str(results_dir / "best_model.safetensors"),
best_optimizer_path=str(results_dir / "best_optimizer.pt"),
best_metrics_path=str(results_dir / "best_metrics.csv"),
best_config_path=str(results_dir / "best_config.yaml"),
debug_mode=True,
)

except RuntimeError as e:
if "zero_division" in str(e).lower():
pytest.skip("Skipping due to known metric edge case")
raise
finally:
# Ensure Ray is shut down properly
if ray.is_initialized():
ray.shutdown()

# Clean up any ray files/directories that may have been created
ray_results_dir = os.path.expanduser("~/ray_results")
ray_results_dir = os.path.expanduser("tests/test_data/titanic/test_results/")
if os.path.exists(ray_results_dir):
import shutil

shutil.rmtree(ray_results_dir)

0 comments on commit 503048a

Please sign in to comment.