Add test coverage for upload_mlflow_artifacts wrapper function#383
Conversation
…ction Co-authored-by: gphuang <13152353+gphuang@users.noreply.github.com>
…rtions Co-authored-by: gphuang <13152353+gphuang@users.noreply.github.com>
Co-authored-by: gphuang <13152353+gphuang@users.noreply.github.com>
gphuang
requested review from
Xiaoming-AMD,
limou102 and
wenxie-amd
as code owners
There was a problem hiding this comment.
Pull request overview
This PR adds comprehensive test coverage for the upload_mlflow_artifacts wrapper function in global_vars.py, which coordinates MLflow artifact uploads but previously lacked any tests.
- Introduces 8 test cases covering parameter forwarding, early returns, and edge cases
- Tests verify correct delegation to
upload_artifacts_to_mlflowwith various parameter combinations - Validates handling of disabled MLflow scenarios (when
mlflow_writeris None)
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| @patch("primus.backends.megatron.training.global_vars.get_mlflow_writer") | ||
| @patch("primus.backends.megatron.training.global_vars.upload_artifacts_to_mlflow") | ||
| def test_none_mlflow_writer_early_return(self, mock_upload_artifacts, mock_get_writer): | ||
| """Test early return when MLflow is not enabled.""" | ||
| # Setup mock to return None | ||
| mock_get_writer.return_value = None | ||
|
|
||
| # Call with all parameters set | ||
| result = upload_mlflow_artifacts( | ||
| tensorboard_dir="/path/to/tensorboard", | ||
| exp_root_path="/path/to/exp", | ||
| upload_traces=True, | ||
| upload_logs=True, | ||
| upload_tracelens_report=True, | ||
| tracelens_ranks=[0], | ||
| tracelens_max_reports=5, | ||
| tracelens_output_format="xlsx", | ||
| ) | ||
|
|
||
| # Verify no artifact upload was attempted | ||
| mock_upload_artifacts.assert_not_called() | ||
|
|
||
| # Verify None is returned immediately | ||
| assert result is None |
There was a problem hiding this comment.
This test is redundant with test_none_mlflow_writer_returns_none (lines 96-114). Both tests verify the same behavior: that the function returns None and doesn't call upload_artifacts_to_mlflow when mlflow_writer is None. The only difference is the parameters passed, but since the early return happens before any parameter processing, passing different parameters doesn't test any additional logic. Consider removing this test to reduce maintenance overhead and improve test clarity.
| mock_upload_artifacts.return_value = {"traces": 0, "logs": 0, "tracelens_reports": 0} | ||
|
|
||
| # Call with empty ranks list | ||
| result = upload_mlflow_artifacts( |
There was a problem hiding this comment.
Variable result is not used.
The
upload_mlflow_artifactsfunction inglobal_vars.pylacked test coverage despite being a coordination wrapper for MLflow artifact uploads.Changes
test_global_vars.pywith 8 test cases covering:upload_artifacts_to_mlflow(defaults and custom values)mlflow_writeris NoneTest Pattern
All tests pass. No security issues detected.
💡 You can make Copilot smarter by setting up custom instructions, customizing its development environment and configuring Model Context Protocol (MCP) servers. Learn more Copilot coding agent tips in the docs.