Skip to content

Commit

Permalink
Add reservations info method (#103)
Browse files Browse the repository at this point in the history 
* Add reservations info method

* Remove testing for deprecated reservation method

* Remove testing for deprecated reservation method

* Fix merge
  • Loading branch information
@ekouts
ekouts authored Apr 29, 2024
1 parent dfc9f1a commit efb8044
Show file tree
Hide file tree
Showing 8 changed files with 275 additions and 513 deletions.
27 changes: 27 additions & 0 deletions firecrest/AsyncClient.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1306,6 +1306,33 @@ async def partitions(
result = await t.poll_task("200")
return result

async def reservations(
self,
machine: str,
reservations: Optional[Sequence[str]] = None,
) -> List[t.ReservationInfo]:
"""Retrieves information about the reservations.
This call uses the `scontrol show reservations` command.
:param machine: the machine name where the scheduler belongs to
:param reservations: specific reservations to query
:calls: GET `/compute/reservations`
GET `/tasks/{taskid}`
.. warning:: This is available only for FirecREST>=1.16.0
"""
params = {}
if reservations:
params["reservations"] = ",".join(reservations)

resp = await self._get_request(
endpoint="/compute/reservations",
additional_headers={"X-Machine-Name": machine},
params=params,
)
json_response = self._json_response([resp], 200)
t = ComputeTask(self, json_response["task_id"], [resp])
result = await t.poll_task("200")
return result

async def cancel(self, machine: str, job_id: str | int) -> str:
"""Cancels running job.
This call uses the `scancel` command.
Expand Down
29 changes: 29 additions & 0 deletions firecrest/BasicClient.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1146,6 +1146,35 @@ def partitions(
)
return result

def reservations(
self,
machine: str,
reservations: Optional[Sequence[str]] = None,
) -> List[t.ReservationInfo]:
"""Retrieves information about the compute reservations.
This call uses the `scontrol show reservations` command.
:param machine: the machine name where the scheduler belongs to
:param nodes: specific reservations to query
:calls: GET `/compute/reservations`
GET `/tasks/{taskid}`
.. warning:: This is available only for FirecREST>=1.16.0
"""
params = {}
if reservations:
params["reservations"] = ",".join(reservations)

resp = self._get_request(
endpoint="/compute/reservations",
additional_headers={"X-Machine-Name": machine},
params=params,
)
self._current_method_requests.append(resp)
json_response = self._json_response(self._current_method_requests, 200)
result = self._poll_tasks(
json_response["task_id"], "200", iter([1, 0.5, 0.25])
)
return result

def cancel(self, machine: str, job_id: str | int) -> str:
"""Cancels running job.
This call uses the `scancel` command.
Expand Down
124 changes: 32 additions & 92 deletions firecrest/cli/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1409,63 +1409,8 @@ def cancel(
raise typer.Exit(code=1)


@reservation_app.command(name='list', rich_help_panel="Reservation commands")
def list_command(
config_from_parent: str = typer.Option(None,
callback=config_parent_load_callback,
is_eager=True,
hidden=True
),
system: str = typer.Option(
..., "-s", "--system", help="The name of the system.", envvar="FIRECREST_SYSTEM"
),
):
"""List all active reservations and their status"""
try:
res = client.all_reservations(system)
console.print(res)
except Exception as e:
examine_exeption(e)
raise typer.Exit(code=1)


@reservation_app.command(rich_help_panel="Reservation commands")
def create(
config_from_parent: str = typer.Option(None,
callback=config_parent_load_callback,
is_eager=True,
hidden=True
),
system: str = typer.Option(
..., "-s", "--system", help="The name of the system.", envvar="FIRECREST_SYSTEM"
),
name: str = typer.Argument(..., help="The reservation name."),
account: str = typer.Argument(
..., help="The account in SLURM to which the reservation is made for."
),
num_nodes: str = typer.Argument(
..., help="The number of nodes needed for the reservation."
),
node_type: str = typer.Argument(..., help="The node type."),
start_time: str = typer.Argument(
..., help="The start time for reservation (YYYY-MM-DDTHH:MM:SS)."
),
end_time: str = typer.Argument(
..., help="The end time for reservation (YYYY-MM-DDTHH:MM:SS)."
),
):
"""Create a reservation"""
try:
client.create_reservation(
system, name, account, num_nodes, node_type, start_time, end_time
)
except Exception as e:
examine_exeption(e)
raise typer.Exit(code=1)


@reservation_app.command(rich_help_panel="Reservation commands")
def update(
@app.command(rich_help_panel="Compute commands")
def get_reservations(
config_from_parent: str = typer.Option(None,
callback=config_parent_load_callback,
is_eager=True,
Expand All @@ -1474,46 +1419,41 @@ def update(
system: str = typer.Option(
..., "-s", "--system", help="The name of the system.", envvar="FIRECREST_SYSTEM"
),
name: str = typer.Argument(..., help="The reservation name."),
account: str = typer.Argument(
..., help="The account in SLURM to which the reservation is made for."
),
num_nodes: str = typer.Argument(
..., help="The number of nodes needed for the reservation."
),
node_type: str = typer.Argument(..., help="The node type."),
start_time: str = typer.Argument(
..., help="The start time for reservation (YYYY-MM-DDTHH:MM:SS)."
),
end_time: str = typer.Argument(
..., help="The end time for reservation (YYYY-MM-DDTHH:MM:SS)."
reservations: Optional[List[str]] = typer.Argument(
None, help="List of specific reservations to query."
),
raw: bool = typer.Option(False, "--raw", help="Print unformatted."),
):
"""Update a reservation"""
"""Retrieves information about the reservations.
This call uses the `scontrol show reservations` command
"""
try:
client.update_reservation(
system, name, account, num_nodes, node_type, start_time, end_time
)
except Exception as e:
examine_exeption(e)
raise typer.Exit(code=1)
results = client.reservations(system, reservations)
if raw:
console.print(results)
else:
parsed_results = []
for item in results:
parsed_item = {}
for key, value in item.items():
if isinstance(value, list):
parsed_item[key] = ", ".join(value)
else:
parsed_item[key] = str(value)

parsed_results.append(parsed_item)

@reservation_app.command(rich_help_panel="Reservation commands")
def delete(
config_from_parent: str = typer.Option(None,
callback=config_parent_load_callback,
is_eager=True,
hidden=True
),
system: str = typer.Option(
..., "-s", "--system", help="The name of the system.", envvar="FIRECREST_SYSTEM"
),
name: str = typer.Argument(..., help="The reservation name."),
):
"""Delete a reservation"""
try:
client.delete_reservation(system, name)
table = create_table(
"Information about reservations in the system",
parsed_results,
("Name", "ReservationName"),
("State", "State"),
("Nodes", "Nodes"),
("StartTime", "StartTime"),
("EndTime", "EndTime"),
("Features", "Features"),
)
console.print(table)
except Exception as e:
examine_exeption(e)
raise typer.Exit(code=1)
Expand Down
10 changes: 10 additions & 0 deletions firecrest/types.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -158,6 +158,16 @@ class PartitionInfo(TypedDict):
TotalNodes: str


class ReservationInfo(TypedDict):
"""A job queue record, from `compute/reservations`"""

ReservationName: str
State: str
Nodes: str
StartTime: str
EndTime: str
Features: str


class JobSubmit(TypedDict):
"""A job submit record, from `compute/jobs`"""
Expand Down
Loading

0 comments on commit efb8044

Please sign in to comment.