Add Reservation

Author: AryanAhadinia

.github/workflows/ci-comprehensive.yml

@@ -23,6 +23,7 @@ jobs:
       FIRECREST_TOKEN_URI: ${{ secrets.FIRECREST_TOKEN_URI }}
       FIRECREST_URL: ${{ secrets.FIRECREST_URL }}
       FIRECREST_USERNAME: ${{ secrets.FIRECREST_USERNAME }}
+      RESERVATION: ${{ secrets.RESERVATION }}
     steps:
       - uses: actions/checkout@v5
       - uses: ./.github/actions/setup

.github/workflows/ci-lightweight.yml

@@ -23,6 +23,7 @@ jobs:
       FIRECREST_TOKEN_URI: ${{ secrets.FIRECREST_TOKEN_URI }}
       FIRECREST_URL: ${{ secrets.FIRECREST_URL }}
       FIRECREST_USERNAME: ${{ secrets.FIRECREST_USERNAME }}
+      RESERVATION: ${{ secrets.RESERVATION }}
     steps:
       - uses: actions/checkout@v5
       - uses: ./.github/actions/setup

README.md

@@ -79,6 +79,7 @@ If you want to skip the interactive prompts and launch a pre-configured model di
 | ----------------------- | -------------------------------- | ---------------------------------------------------------------------- |
 | `--firecrest-system`    | `SML_FIRECREST_SYSTEM`           | Target system to launch on (required only if using FirecREST launcher) |
 | `--partition`           | `SML_PARTITION`                  | SLURM partition to use                                                 |
+| `--reservation`         | `SML_RESERVATION`                | SLURM reservation name (optional)                                      |
 | `--model`               |                                  | Model to launch (`<vendor>/<model>`)                                   |
 | `--framework`           |                                  | Inference framework to use                                             |
 | `--workers`             |                                  | Number of workers                                                      |
@@ -117,6 +118,7 @@ For full control over the SLURM job, use `sml advanced`. This bypasses the model
 | -------------------------- | ------------------------- | ----------------------------------------------------------------- |
 | `--firecrest-system`       | `SML_FIRECREST_SYSTEM`    | Target HPC system to launch on                                    |
 | `--partition`              | `SML_PARTITION`           | SLURM partition to use                                            |
+| `--slurm-reservation`      | `SML_RESERVATION`         | SLURM reservation name (optional)                                 |
 | `--serving-framework`      |                           | Inference framework (`sglang`, `vllm`) — **required**             |
 | `--slurm-environment`      |                           | Local path to the environment `.toml` file — **required**         |
 | `--framework-args`         |                           | Arguments forwarded to the inference framework                    |
@@ -174,6 +176,7 @@ export FIRECREST_SYSTEM=clariden
 export FIRECREST_ACCOUNT=<your-account>
 export FIRECREST_PARTITION=normal
 export CSCS_API_KEY=<your-api-key>
+export RESERVATION=<your-reservation>
 ```
 
 This file will be sourced when running the tests with `make test`, and the environment variables will be available for the tests.

src/swiss_ai_model_launch/assets/template.jinja

@@ -5,6 +5,7 @@
 #SBATCH --exclusive
 #SBATCH --nodes={{ nodes }}
 #SBATCH --partition={{ partition }}
+{% if reservation %}#SBATCH --reservation={{ reservation }}{% endif %}
 #SBATCH --output=logs/%j/log.out
 #SBATCH --error=logs/%j/log.out
 

src/swiss_ai_model_launch/cli/main.py

@@ -73,6 +73,19 @@ def _make_partition_config(
     )
 
 
+def _make_reservation_config() -> ChainConfiguration:
+    return ChainConfiguration(
+        name="reservation_configuration",
+        chain=[
+            TextConfiguration(
+                name="reservation",
+                prompt="SLURM reservation name (optional, leave blank to skip).",
+                env_var="SML_RESERVATION",
+            ),
+        ],
+    )
+
+
 def _make_launch_request_config(
     vendor_models_factory: _OptionsFactory = None,
     frameworks_factory: _OptionsFactory = None,
@@ -148,6 +161,7 @@ def _build_parser() -> argparse.ArgumentParser:
     )
     _make_firecrest_launcher_config().add_to_parser(preconfigured_parser)
     _make_partition_config().add_to_parser(preconfigured_parser)
+    _make_reservation_config().add_to_parser(preconfigured_parser)
     _make_launch_request_config().add_to_parser(preconfigured_parser)
 
     advanced_parser = subparsers.add_parser(
@@ -203,6 +217,13 @@ def _build_parser() -> argparse.ArgumentParser:
         metavar="HH:MM:SS",
         help="Job time limit (default: 00:05:00).",
     )
+    advanced_parser.add_argument(
+        "--slurm-reservation",
+        dest="reservation",
+        default=None,
+        metavar="RESERVATION",
+        help="SLURM reservation name (optional).",
+    )
     advanced_parser.add_argument(
         "--served-model-name",
         dest="served_model_name",
@@ -290,12 +311,16 @@ async def _get_partitions() -> dict[str, tuple[str, str]]:
     partition_config = _make_partition_config(partitions_factory=_get_partitions)
     await partition_config.aconfigure(args=args)
 
+    reservation_config = _make_reservation_config()
+    await reservation_config.aconfigure(args=args)
+
     return FirecRESTLauncher(
         client,
         system_name=system_name,
         username=user_info["user"]["name"],
         account=user_info["group"]["name"],
         partition=partition_config.get_non_none_value("partition"),
+        reservation=reservation_config.get_value("reservation") or None,
         telemetry_endpoint=telemetry_endpoint,
     )
 
@@ -320,11 +345,15 @@ async def _get_partitions() -> dict[str, tuple[str, str]]:
     partition_config = _make_partition_config(partitions_factory=_get_partitions)
     await partition_config.aconfigure(args=args)
 
+    reservation_config = _make_reservation_config()
+    await reservation_config.aconfigure(args=args)
+
     return SlurmLauncher(
         system_name="local",
         username=getpass.getuser(),
         account=grp.getgrgid(os.getgid()).gr_name,
         partition=partition_config.get_non_none_value("partition"),
+        reservation=reservation_config.get_value("reservation") or None,
         telemetry_endpoint=telemetry_endpoint,
     )
 
@@ -541,6 +570,7 @@ async def _run_advanced(args: argparse.Namespace) -> None:
         nodes_per_worker=args.nodes_per_worker,
         nodes=args.nodes,
         time=args.time,
+        reservation=args.reservation or None,
         environment=args.slurm_environment,
         framework=args.framework,
         framework_args=args.framework_args,

src/swiss_ai_model_launch/launchers/firecrest_launcher.py

@@ -29,13 +29,15 @@ def __init__(
         username: str,
         account: str,
         partition: str,
+        reservation: str | None = None,
         telemetry_endpoint: str | None = None,
     ):
         super().__init__(
             system_name=system_name,
             username=username,
             account=account,
             partition=partition,
+            reservation=reservation,
             telemetry_endpoint=telemetry_endpoint,
         )
         self.client = client
@@ -64,6 +66,7 @@ def _get_launch_args_from_request(
             workers=launch_request.workers,
             nodes_per_worker=launch_request.nodes_per_worker,
             time=launch_request.time,
+            reservation=self.reservation,
             environment=launch_request.environment,
             framework=launch_request.framework,
             served_model_name=served_model_name,
@@ -128,7 +131,9 @@ async def launch_with_args(self, launch_args: LaunchArgs) -> tuple[int, str]:
         remote_env_path = await self._upload_env_file(
             launch_args.environment, launch_args.framework
         )
-        launch_args = launch_args.model_copy(update={"environment": remote_env_path})
+        launch_args = launch_args.model_copy(
+            update={"environment": remote_env_path, "reservation": self.reservation}
+        )
         script_str = render_job_script(launch_args)
         job_submission_report = await self.client.submit(
             system_name=self.system_name,

src/swiss_ai_model_launch/launchers/launch_args.py

@@ -12,6 +12,7 @@ class LaunchArgs(BaseModel):
     nodes: int | None = None
 
     time: str = "00:05:00"
+    reservation: str | None = None
     environment: str
 
     framework: str

src/swiss_ai_model_launch/launchers/launcher.py

@@ -26,12 +26,14 @@ def __init__(
         username: str,
         account: str,
         partition: str,
+        reservation: str | None = None,
         telemetry_endpoint: str | None = None,
     ):
         self.system_name = system_name
         self.username = username
         self.account = account
         self.partition = partition
+        self.reservation = reservation
         self.telemetry_endpoint = telemetry_endpoint
 
     @abstractmethod

src/swiss_ai_model_launch/launchers/slurm_launcher.py

@@ -27,6 +27,7 @@ def __init__(
         username: str,
         account: str,
         partition: str,
+        reservation: str | None = None,
         model_registry: Path = _REMOTE_MODEL_REGISTRY,
         telemetry_endpoint: str | None = None,
     ):
@@ -35,6 +36,7 @@ def __init__(
             username=username,
             account=account,
             partition=partition,
+            reservation=reservation,
             telemetry_endpoint=telemetry_endpoint,
         )
         self.model_registry = model_registry
@@ -59,6 +61,7 @@ def _get_launch_args_from_request(
             workers=launch_request.workers,
             nodes_per_worker=launch_request.nodes_per_worker,
             time=launch_request.time,
+            reservation=self.reservation,
             environment=launch_request.environment,
             framework=launch_request.framework,
             served_model_name=served_model_name,
@@ -138,6 +141,7 @@ async def get_preconfigured_models(self) -> list[LaunchRequest]:
         ]
 
     async def launch_with_args(self, launch_args: LaunchArgs) -> tuple[int, str]:
+        launch_args = launch_args.model_copy(update={"reservation": self.reservation})
         job_id = await self._sbatch(launch_args)
         return job_id, launch_args.served_model_name
 

tests/integration/test_firecrest_launcher.py

@@ -38,6 +38,7 @@
     "FIRECREST_ACCOUNT",
     "FIRECREST_PARTITION",
     "CSCS_API_KEY",
+    "RESERVATION",
 ]
 
 
@@ -68,6 +69,7 @@ def launcher(env: dict[str, str]) -> FirecRESTLauncher:
         username=env["FIRECREST_USERNAME"],
         account=env["FIRECREST_ACCOUNT"],
         partition=env["FIRECREST_PARTITION"],
+        reservation=env["RESERVATION"] or None,
     )