Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add reservations info method #103

Merged
merged 5 commits into from
Apr 29, 2024
Merged

Add reservations info method #103

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
7758c75
Add reservations info method
ekouts Apr 28, 2024
9232869
Remove testing for deprecated reservation method
ekouts Apr 29, 2024
dc6a346
Remove testing for deprecated reservation method
ekouts Apr 29, 2024
da49ad7
Merge branch 'main' of github.com:eth-cscs/pyfirecrest into new_metho…
ekouts Apr 29, 2024
c690706
Fix merge
ekouts Apr 29, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 27 additions & 0 deletions firecrest/AsyncClient.py
View file
Open in desktop
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
Open in desktop
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
Open in desktop
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
Open in desktop
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
Loading