feat: add tool to check if GPUs are available for training #102
Conversation
vaibs-d
changed the title
![cubic-dev-ai[bot]](https://avatars.githubusercontent.com/in/1082092?s=60&v=4){kind=small}
There was a problem hiding this comment.
Nice feature @vaibs-d ! I have a couple of questions about how the new tools plug into the agentic workflow, see comments. Otherwise, excited to get this feature merged 🚀
| def get_executor_tool(distributed: bool = False) -> Callable: | ||
| """Get the appropriate executor tool based on the distributed flag.""" | ||
| def get_executor_tool() -> Callable: | ||
| """Get the executor tool for training code execution.""" |
There was a problem hiding this comment.
We actually don't need this to be wrapped in a closure anymore, since we're not passing any parameters to it.
| @tool | ||
| def get_gpu_info() -> dict: | ||
| """ | ||
| Get available GPU information for code generation. |
There was a problem hiding this comment.
I have two questions about this tool:
- It doesn't seem to be added to the list of tools for any of the agents, unless I'm missing something?
- I'm not sure how the agent is expected to use this tool. Is the idea that the agent would check using this tool if different ML frameworks have access to GPU?
I think I see a couple of issues here:
- The docstring may be a little too terse for an agent to reliably understand what this is for.
- This should be added to an agent's tools, and the agent's prompts should maybe also be modified to suggest using this
- This will check for GPU availability on the local compute instance; if we're using a Ray cluster, this would not check for the presence of GPUs in the Ray cluster
| @tool | ||
| def get_ray_info() -> dict: | ||
| """ | ||
| Get Ray cluster information including GPU availability. |
There was a problem hiding this comment.
Same as the tool above, I'm not entirely clear on how an agent should use this tool. The docstring probably doesn't give an agent enough context about what this tool is really "for" and how it fits in the broader workflow.
At least with GPT-4o and Claude 3.7 I've seen that the tool docstrings need to be very unambiguous for the agent to reliably use the tool correctly.
vaibs-d
force-pushed
the
add-gpu-support
branch
from
0e5c7db to
903b1b9
Compare
This PR adds tools for the MLE agent to asses if GPUs are available via Ray or locally to use for model training.