|
1 | 1 | # encoding: utf-8
|
2 | 2 | from __future__ import absolute_import
|
3 | 3 |
|
| 4 | +import time |
| 5 | +import typing |
4 | 6 | from enum import Enum
|
5 | 7 |
|
| 8 | +import requests |
| 9 | + |
6 | 10 | from rapyuta_io.clients.device import Device, DeviceStatus
|
| 11 | +from rapyuta_io.clients.model import Command |
7 | 12 | from rapyuta_io.utils import RestClient
|
| 13 | +from rapyuta_io.utils.error import ParameterMissingException |
8 | 14 | from rapyuta_io.utils.rest_client import HttpMethod
|
9 |
| -from rapyuta_io.utils.settings import DEVICE_API_PATH, DEVICE_SELECTION_API_PATH, PARAMETERS_API_PATH, \ |
10 |
| - DEVICE_API_ADD_DEVICE_PATH, DAEMONS_PATH |
11 |
| -from rapyuta_io.utils.utils import create_auth_header, prepend_bearer_to_auth_token, get_api_response_data, \ |
| 15 | +from rapyuta_io.utils.settings import DAEMONS_PATH, DEVICE_API_ADD_DEVICE_PATH, DEVICE_API_PATH, \ |
| 16 | + DEVICE_COMMAND_API_PATH, DEVICE_SELECTION_API_PATH, PARAMETERS_API_PATH |
| 17 | +from rapyuta_io.utils.utils import create_auth_header, get_api_response_data, get_error, prepend_bearer_to_auth_token, \ |
12 | 18 | validate_list_of_strings
|
13 | 19 |
|
14 | 20 |
|
@@ -131,3 +137,65 @@ def patch_daemons(self, device_id, payload):
|
131 | 137 | headers = create_auth_header(self._auth_token, self._project)
|
132 | 138 | response = RestClient(url).method(HttpMethod.PATCH).headers(headers).execute(payload=payload)
|
133 | 139 | return get_api_response_data(response, parse_full=True)
|
| 140 | + |
| 141 | + def execute_command( |
| 142 | + self, |
| 143 | + device_ids: typing.List[str], |
| 144 | + command: Command, |
| 145 | + retry_limit: int = 0, |
| 146 | + retry_interval: int = 10, |
| 147 | + timeout: int = 300, |
| 148 | + ): |
| 149 | + """Execute a command on the specified devices. |
| 150 | +
|
| 151 | + Args: |
| 152 | + device_ids: List of device IDs on which the command should be executed. |
| 153 | + command: Command object to be executed. |
| 154 | + retry_limit: Number of retries in case of API failure. |
| 155 | + retry_interval: Interval between retries. |
| 156 | + timeout: Maximum time to wait for the background command to finish. |
| 157 | +
|
| 158 | + Returns: |
| 159 | + dict: Output of the command execution. |
| 160 | +
|
| 161 | + Raises: |
| 162 | + ValueError: If device_ids is empty. |
| 163 | + TimeoutError: If command execution takes longer than the specified timeout. |
| 164 | + ParameterMissingException: If the command is missing required parameters. |
| 165 | + """ |
| 166 | + if not device_ids: |
| 167 | + raise ValueError("device_ids cannot be empty") |
| 168 | + |
| 169 | + command.validate() |
| 170 | + command.device_ids = device_ids |
| 171 | + |
| 172 | + url = self._device_api_host + DEVICE_COMMAND_API_PATH |
| 173 | + rc = RestClient(url).method(HttpMethod.POST).headers( |
| 174 | + create_auth_header(self._auth_token, self._project)) |
| 175 | + response = rc.retry(retry_limit).execute(payload=command.to_json()) |
| 176 | + if response.status_code == requests.codes.BAD_REQUEST: |
| 177 | + raise ParameterMissingException(get_error(response.text)) |
| 178 | + |
| 179 | + execution_result = get_api_response_data(response) |
| 180 | + |
| 181 | + if not command.bg: |
| 182 | + return execution_result |
| 183 | + |
| 184 | + jid = execution_result.get('jid') |
| 185 | + if not jid: |
| 186 | + raise ValueError("job id not found in the response") |
| 187 | + |
| 188 | + url = self._device_api_host + DEVICE_COMMAND_API_PATH + jid |
| 189 | + query = {"jid": jid, "device_id": device_ids} |
| 190 | + time_elapsed = 0 |
| 191 | + wait_interval = retry_interval |
| 192 | + while time_elapsed < timeout: |
| 193 | + response = RestClient(url).method(HttpMethod.GET).headers( |
| 194 | + create_auth_header(self._auth_token, self._project)).query_param(query_param=query).execute() |
| 195 | + if response.status_code == requests.codes.OK: |
| 196 | + result = get_api_response_data(response) |
| 197 | + return result |
| 198 | + time.sleep(wait_interval) |
| 199 | + time_elapsed += wait_interval |
| 200 | + |
| 201 | + raise TimeoutError(f"command result not available after {timeout} seconds") |
0 commit comments