fix(scheduler): Fix SLURM scheduler spec.image and spec.additional_bash_cmds arguments.

[nuzant]

Description

Fix issues in the Slurm scheduler:

1. Remove hardcoded container_image parameter - use spec.image from scheduling spec instead
2. Properly execute additional bash commands before the RPC server, specified by spec.additional_bash_cmds arguments.
3. Fix slurm scheduler tests.

Related Issue

N/A

Type of Change

• Bug fix (non-breaking change that fixes an issue)
• New feature (non-breaking change that adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Documentation update
• Code refactoring (no functional changes)
• Performance improvement
• Test coverage improvement

Checklist

• I have read the Contributing Guide
• I have run formatting tools (pre-commit or manual)
• I have run relevant unit tests and they pass
• I have added tests for new functionality
• I have updated documentation if needed
• My branch is up to date with main
• This PR introduces breaking changes (if yes, fill out details below)
• If this PR changes documentation, I have built and previewed it locally with jb build docs
• No critical issues raised by AI reviewers (/gemini review)

Breaking Change Details (if applicable):

N/A

Additional Context

Generated with Claude Code

[gemini-code-assist]

Summary of Changes

Hello @nuzant, 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!

This pull request significantly enhances the Slurm scheduler by addressing key issues related to container image management and command execution. It eliminates a hardcoded container image, allowing for dynamic image specification, and rectifies a bug where native mode failed to execute additional bash commands. These changes ensure more consistent and robust worker provisioning, particularly for scenarios involving forked processes, and improve the overall flexibility of the Slurm integration.

Highlights

• Slurm Scheduler Container Image: The hardcoded container image path in the Slurm scheduler has been removed. It now dynamically uses the spec.image from the scheduling specification, providing greater flexibility.
• Native Mode Bash Command Execution: The Slurm scheduler's native (non-Apptainer) mode now correctly executes all additional_bash_cmds specified in the scheduling specification before initiating the RPC server command. Previously, these commands were ignored.
• Consistent Bash Command Handling: Bash command execution logic has been standardized across both Apptainer and native container modes within the Slurm scheduler, ensuring consistent behavior for additional_bash_cmds.
• Dynamic Port Allocation API: A new FastAPI endpoint /alloc_ports has been introduced in the OpenAI proxy server. This endpoint allows for the dynamic allocation of multiple free ports, which is crucial for supporting the SlurmScheduler's fork_workers mechanism.
• Updated Slurm Scheduler Tests: Tests for the Slurm scheduler have been updated to accurately reflect the new behavior of fork_workers, which now creates distinct worker processes with unique IDs and ports, rather than reusing target worker information.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

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 by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

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 pull request 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 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. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

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.

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 several fixes and improvements to the Slurm scheduler. It removes a hardcoded container image, making it configurable via the scheduling spec. It also fixes an issue where native mode (non-apptainer) would ignore additional bash commands, and unifies command execution between native and apptainer modes. Additionally, it adds port allocation capabilities to the proxy server to support forked workers.

My review focuses on two main areas. First, I've identified a potential race condition in the new port allocation endpoint that could lead to the same port being assigned twice. I've suggested using an asyncio.Lock to ensure thread safety. Second, I've found a bug in how shell commands are escaped when constructing the sbatch script, which could cause issues with commands containing special characters. I've recommended using shlex.quote for a more robust solution.

The rest of the changes, including the test updates for forked workers, look good and align with the goals of the pull request.

[garrett4wade]

The change is not necessarily related to CI. Merging anyways.