Refactor MLflow client usage in tests to streamline configuration loading #97
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
…nto a new module. This change consolidates various configuration definitions previously scattered across multiple files, improving organization and maintainability. The new module includes configurations for optimizers, logging, persistence, generative processes, predictive models, training, validation, and MLflow, along with their respective validation functions.
… into a dedicated module. This update removes obsolete config files and centralizes the configuration for generative processes, logging, persistence, predictive models, and training, enhancing organization and maintainability. Additionally, update imports across the codebase to reflect the new structure.
… ensure nested configuration integrity.
…c config classes, enhancing flexibility and consistency across configuration validation. Update validation logic to ensure required fields are checked and improve error messaging for better clarity.
…or logging, generative processes, optimizers, and persistence. This update improves the robustness of configuration management by ensuring that each component's configuration is validated before being processed, leading to better error handling and logging.
…resolution logging for missing configuration fields. This update improves error handling by ensuring that missing values are logged and only validated when present, maintaining the integrity of generative process configurations.
…nhancing flexibility in configuration. Update validation logic to check URIs only when provided, improving error handling and maintaining configuration integrity.
…proved error handling and logging. This change enhances clarity in error messages during configuration validation.
…y raise ConfigValidationError for all validation failures, improving error handling and clarity in configuration validation across various components.
…r configuration management, enhancing flexibility and consistency. Implement validation checks for required fields, ensuring robust error handling with ConfigValidationError for missing parameters.
…figs.py This commit introduces a new test suite for validating logger targets and configurations using pytest. It includes tests for valid and invalid logger targets, configurations, and validation logic, ensuring robust error handling with ConfigValidationError. The tests cover various scenarios, including missing fields and incorrect types, enhancing the reliability of the logging configuration management.
… methods for improved encapsulation. This change enhances code clarity and maintains consistent error handling with ConfigValidationError across various configuration validations.
…ault value is used when not explicitly provided. This change improves configuration flexibility and maintains consistency in registry URI resolution.
…hancing configuration flexibility. Adjust validation logic to accommodate the new type, ensuring robust error handling for optional boolean values.
…_configs.py This commit introduces tests for the validate_mlflow_config function, covering both valid and invalid configurations. The tests ensure that required fields are present and validate error handling with ConfigValidationError for missing parameters, enhancing the reliability of MLFlow configuration management.
…t_structured_configs.py This commit introduces a comprehensive suite of tests for validating generative process configurations and targets. It includes tests for both valid and invalid configurations, ensuring that the validation functions correctly raise ConfigValidationError for various error scenarios. This enhancement improves the reliability and robustness of generative process configuration management.
…figs.py This commit updates the validation logic to ensure that the vocab_size is positive and adds checks for base_vocab_size and special tokens (bos_token, eos_token) to prevent invalid configurations. Additionally, new unit tests are introduced to cover various invalid scenarios, ensuring robust error handling with ConfigValidationError. These changes improve the reliability of generative process configuration management.
…ed validation functions for positive and non-negative integers. This change enhances code clarity and reduces redundancy in the validation process for various configuration parameters, ensuring consistent error handling with ConfigValidationError.
…alues for optional positive and non-negative integers. This update improves error handling by ensuring that configurations can accept None where appropriate, while maintaining consistent validation logic. Additionally, unit tests in test_structured_configs.py are updated to reflect these changes, ensuring comprehensive coverage for various scenarios.
…nction for non-empty string checks. This change enhances consistency across various configuration validations by replacing the previous optional name validation with a more robust method. Additionally, update unit tests in test_structured_configs.py to reflect the new validation logic, ensuring accurate error messages for invalid name scenarios.
This commit introduces a new feature that allows users to bootstrap configurations by loading sections from previous MLflow runs. The `load_configs` block can be added to the configuration files, enabling the retrieval of specific model and process configurations from past experiments. The implementation includes a new `_load_config` function in `run_management.py` that handles the loading and merging of configurations, along with corresponding unit tests to ensure functionality and correctness.
…ding This commit replaces the stubbed MLflow client with a new helper function, `_create_run_with_config`, to manage MLflow run creation and configuration loading in tests. The updated tests now directly utilize the MLflow API for logging artifacts, enhancing clarity and reducing complexity in the test setup. This change improves the maintainability of the test suite while ensuring accurate simulation of MLflow interactions.
This commit introduces a structured approach to testing MLFlow, logger, and generative process configurations by organizing tests into dedicated classes. It improves the clarity and maintainability of the test suite while ensuring comprehensive coverage for valid and invalid configurations. Additionally, the validation logic is reinforced to handle various edge cases, enhancing error handling with ConfigValidationError across different configuration scenarios.
…ation checks This commit introduces two new configuration parameters, sequence_len and batch_size, to the GenerativeProcessConfig class in structured_configs.py. Validation logic is added to ensure these parameters are positive integers or None, enhancing the robustness of configuration management. Corresponding unit tests are also added in test_structured_configs.py to cover various scenarios, including invalid inputs for both parameters, ensuring comprehensive error handling with ConfigValidationError.
… sequence_len This commit adds the batch_size and sequence_len parameters to the mess3.yaml configuration file, enhancing the generative process setup. Additionally, the _create_initial_state function in run_management.py is refactored to accept batch_size as an argument, allowing for more flexible initial state creation based on the provided configuration. This change improves the adaptability of the generative process management.
…d attribute handling This commit modifies the generative_process and predictive_model configuration files to set default values for sequence_len and n_ctx, enhancing the clarity of model configurations. Additionally, the run_management.py file is refactored to introduce a new function for retrieving attribute values, improving code reusability and maintainability. The changes ensure that configurations are more robust and adaptable to various model setups.
This commit introduces a new BaseConfig class and a corresponding validate_base_config function in structured_configs.py to validate configuration settings such as seed, tags, and mlflow. The validation logic ensures that seed is a non-negative integer, tags are a dictionary with string keys and values, and mlflow adheres to the MLFlowConfig structure. Additionally, comprehensive unit tests are added in test_structured_configs.py to verify the correctness of the validation logic, covering various valid and invalid scenarios to enhance error handling with ConfigValidationError.
This commit updates the HookedTransformerConfigConfig class by reorganizing its attributes and setting default values for several fields, including n_heads and d_vocab. The validation function is enhanced to ensure proper checks for n_heads when it is set to -1, ensuring that d_model is divisible by d_head. Additionally, new tests are added to validate configurations with minimal and full setups, as well as to check for incompatible dimensions, improving overall robustness and error handling in configuration management.
This commit refactors the run_management.py file to enhance the handling of MLFlow and generative process configurations. The changes include replacing specific configuration classes with DictConfig for greater flexibility, improving the retrieval of configuration attributes, and ensuring that required fields like experiment_name and run_name are validated. Additionally, the logic for resolving generative process and hooked transformer configurations is updated to streamline the setup process and improve code maintainability.
This commit updates the HookedTransformerConfig class to include a default _target_ attribute and refines its initialization. Additionally, new tests are introduced for GenerativeProcessConfig, MLFlowConfig, LoggingConfig, and PersistenceConfig, ensuring proper creation and validation of these configurations. The test suite is expanded to cover various scenarios, improving overall test coverage and robustness in configuration management.
…tests This commit modifies the error messages in the tests for various configuration classes to specify that the _target_ attribute cannot be empty or whitespace. The changes enhance clarity in error reporting during validation, ensuring that users receive more informative feedback when encountering configuration issues.
Code Review for PR #97Thank you for this contribution! The config loading feature is a useful addition. Below is a detailed review. Critical Issues1. Feature Not Actually Used (High Priority)The new _get_config() function is defined but never called. The managed_run decorator at line 599 still uses get_config() directly. Impact: The entire load_configs feature won't work until this is fixed. 2. SQL Injection Vulnerability (Security)Line 129 constructs MLflow filter strings with unsanitized user input: Risk: Special characters in run_name could break queries or enable injection attacks. Code Quality Issues3. Missing Error Handling
4. Type Annotation IssuesLine 147 has type: ignore which suppresses type checking. Per CLAUDE.md, code must pass pyright in standard mode. 5. Incomplete DocstringsBoth _load_config and _get_config have minimal docstrings that don't follow Google style required by CLAUDE.md. Test Coverage Issues6. Insufficient Test CoverageTests only cover happy paths. Missing:
Add negative test cases to meet 80% diff-cover requirement. Performance Considerations7. Repeated MLflow Client CreationNew MlflowClient created for each load_config entry. Cache by tracking URI for efficiency. 8. Synchronous DownloadsAll downloads happen sequentially. Consider parallel downloads for multiple entries. Documentation Issues9. Demo Config Invalid ReferencesExample config references non-existent runs. Use placeholders with explanatory comments instead. Positive Aspects
Summary ChecklistMust Fix Before Merge:
Should Fix:
Nice to Have:
Overall AssessmentValuable feature with solid architecture, but needs critical fixes:
Once resolved, this will be a great addition! |
…identifiers. Update README and demo configuration examples to reflect optional use of experiment_id and run_id. Improve error handling for mismatched identifiers in run management tests.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This commit replaces the stubbed MLflow client with a new helper function,
_create_run_with_config, to manage MLflow run creation and configuration loading in tests. The updated tests now directly utilize the MLflow API for logging artifacts, enhancing clarity and reducing complexity in the test setup. This change improves the maintainability of the test suite while ensuring accurate simulation of MLflow interactions.