Skip to content

Commit db34d48

Browse files
feat: add support for background command executes
* chore: add support for bg as true * chore: add timeout support in command
1 parent 3f379ee commit db34d48

File tree

2 files changed

+43
-2
lines changed

2 files changed

+43
-2
lines changed

rapyuta_io/clients/device.py

Lines changed: 39 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import subprocess
55

66
from six.moves.urllib.parse import urlencode
7+
from time import sleep
78
import enum
89

910
import requests
@@ -567,7 +568,44 @@ def execute_command(self, command, retry_limit=0):
567568
if response.status_code == requests.codes.BAD_REQUEST:
568569
raise ParameterMissingException(get_error(response.text))
569570
execution_result = get_api_response_data(response)
570-
return execution_result[self.uuid]
571+
if not command.bg:
572+
return execution_result[self.uuid]
573+
jid = execution_result.get('jid')
574+
if not jid:
575+
raise ValueError("Job ID not found in the response")
576+
return self.fetch_command_result(jid, [self.uuid], timeout=command.timeout)
577+
578+
def fetch_command_result(self, jid: str, deviceids: list, timeout: int):
579+
"""
580+
Fetch the result of the command execution using the job ID (jid) and the first device ID from the list.
581+
Args:
582+
jobid (str): The job ID of the executed command.
583+
deviceids (list): A list of device IDs on which the command was executed.
584+
timeout (int): The maximum time to wait for the result (in seconds). Default is 300 seconds.
585+
Returns:
586+
dict: The result of the command execution.
587+
Raises:
588+
TimeoutError: If the result is not available within the timeout period.
589+
APIError: If the API returns an error.
590+
"""
591+
592+
if not deviceids or not isinstance(deviceids, list):
593+
raise ValueError("Device IDs must be provided as a non-empty list.")
594+
url = self._device_api_host + DEVICE_COMMAND_API_PATH + "jobid"
595+
payload = {
596+
"jid": jid,
597+
"device_id": deviceids[0]
598+
}
599+
total_time_waited = 0
600+
wait_interval = 10
601+
while total_time_waited < timeout:
602+
response = self._execute_api(url, HttpMethod.POST, payload)
603+
if response.status_code == requests.codes.OK:
604+
result = get_api_response_data(response)
605+
return result[deviceids[0]]
606+
sleep(wait_interval)
607+
total_time_waited += wait_interval
608+
raise TimeoutError(f"Command result not available after {timeout} seconds")
571609

572610
def get_config_variables(self):
573611
"""

rapyuta_io/clients/model.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,7 @@ class Command(ObjDict):
6161
6262
"""
6363

64-
def __init__(self, cmd, shell=None, env=None, bg=False, runas=None, pwd=None, cwd=None):
64+
def __init__(self, cmd, shell=None, env=None, bg=False, runas=None, pwd=None, cwd=None, timeout=300):
6565
super(ObjDict, self).__init__()
6666
if env is None:
6767
env = dict()
@@ -73,6 +73,7 @@ def __init__(self, cmd, shell=None, env=None, bg=False, runas=None, pwd=None, cw
7373
self.cwd = pwd
7474
if cwd is not None:
7575
self.cwd = cwd
76+
self.timeout = timeout
7677
self.validate()
7778

7879
def validate(self):
@@ -93,6 +94,8 @@ def validate(self):
9394
raise InvalidCommandException('Invalid environment variables')
9495
return
9596
raise InvalidCommandException('Invalid environment variables')
97+
if self.timeout <= 0:
98+
raise InvalidCommandException("Invalid timeout value")
9699

97100
def to_json(self):
98101
# TODO: we need to rewrite this function.

0 commit comments

Comments
 (0)