feat: Optimisation

[toby-coleman]

Summary

Closes #104 with a Tuner class to allow optimisation using Ray Tune.

Changes

• Tuner class, which performs the optimisation.
• Associated schemas.
• New CLI command: plugboard process tune.

[gemini-code-assist]

Hello @toby-coleman, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

Summary of Changes

This pull request introduces a Tuner class to enable optimization of Plugboard processes using Ray Tune. It addresses issue #104. The changes include the Tuner class itself, associated schemas for specifying the optimization configuration, and updates to the ProcessConfigSpec to include an optional TuneSpec.

Highlights

• Tuner Class: Introduces a Tuner class for performing optimization of Plugboard processes using Ray Tune.
• Schemas: Adds schemas for configuring the optimization process, including TuneSpec, ObjectiveSpec, ParameterSpec, and OptunaSpec.
• Process Configuration: Modifies the ProcessConfigSpec to include an optional TuneSpec for specifying an optimization configuration.

Changelog

Click here to see the changelog

• plugboard/schemas/init.py
  • Added imports for new tune-related schemas: Direction, ObjectiveSpec, OptunaSpec, ParameterSpec, TuneArgsDict, TuneArgsSpec, and TuneSpec.
  • Added the new schemas to the __all__ list.
• plugboard/schemas/_common.py
  • Added validate_assignment=True to the model_config to enable validation on assignment.
  • Added override method to PlugboardBaseModel.
• plugboard/schemas/config.py
  • Imported TuneSpec.
  • Added an optional tune attribute of type TuneSpec | None to the ProcessConfigSpec.
• plugboard/schemas/tune.py
  • Created a new file defining schemas for configuring optimization jobs, including OptunaSpec, BaseFieldSpec, ObjectiveSpec, FloatParameterSpec, IntParameterSpec, CategoricalParameterSpec, ParameterSpec, TuneArgsDict, TuneArgsSpec, and TuneSpec.
• plugboard/tune/init.py
  • Created a new file to define the tune submodule.
  • Added Tuner to __all__.
• plugboard/tune/tune.py
  • Created a new file defining the Tuner class for running optimization on Plugboard processes using Ray Tune.
• pyproject.toml
  • Added optuna>=4.0,<5 to the ray dependencies.
  • Added tune to the ray extra dependencies.
  • Updated openai version from 1.72.0 to 1.73.0.
  • Added alembic, colorlog, and tensorboardx dependencies.
• tests/integration/test_tuner.py
  • Created a new file to define integration tests for the Tuner class.
• uv.lock
  • Updated the lock file to reflect the new dependencies.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

---

A tuner's quest, with Ray's bright aid,
To find the best, a process swayed.
With Optuna's art,
It plays its part,
In Plugboard's tune, a change is made.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces a Tuner class to enable optimization using Ray Tune. The changes include the Tuner class itself and associated schemas. Overall, the code appears well-structured and addresses the issue of optimization effectively. However, there are some areas that could benefit from further attention, particularly regarding error handling, documentation, and code clarity.

Summary of Findings

• Missing Error Handling: The code lacks comprehensive error handling, especially around external library calls (e.g., Ray Tune functions). This could lead to unexpected crashes or incorrect behavior if these calls fail. Additional error handling is needed to ensure the code is robust and reliable.
• TODOs in Code: Several TODO comments exist in the code, indicating incomplete or potentially problematic areas. These need to be addressed to ensure the code is production-ready. Specifically, the TODO regarding ComponentRegistry access on Ray workers needs to be resolved.
• Incomplete Documentation: While the code includes docstrings, some areas, particularly within the Tuner class, could benefit from more detailed explanations. This would improve the code's maintainability and make it easier for others to understand and use.

Merge Readiness

The pull request introduces valuable functionality for optimization. However, before merging, it's crucial to address the identified issues, particularly the missing error handling and TODO comments. Addressing these points will significantly improve the code's robustness and maintainability. I am unable to approve this pull request, and recommend that it not be merged until the critical and high severity issues are addressed (at a minimum), and that others review and approve this code before merging.

[toby-coleman]

@chrisk314 - ready for review.

One issue of note is the need to set the namespace when setting up the RayProcess. This is to prevent actor name clashes when running the same process in parallel. Perhaps overlaps with the job ID - if each Process had its own job ID, then they could each have their own namespace?

[chrisk314]

Looks great. One issue with the schemas which looks like a bug, and one typo spotted; otherwise, just a few suggestions for possible minor code improvements.