diff --git a/README.md b/README.md index 9ab4a4a..e005819 100644 --- a/README.md +++ b/README.md @@ -79,6 +79,9 @@ The arguments for the `WherobotsRunOperator` constructor: If `False`, the operator will not poll the logs, just track the status of the run. * `polling_interval`: The interval in seconds to poll the status of the run. The default value is `30`. +* `timeout_seconds: int`: This parameter sets a maximum run time (in seconds) to prevent runaway processes. +If the specified value exceeds the Max Workload Alive Hours, the timeout will be capped at the maximum permissible limit. +Defaults to `3600` seconds (1 hour). * `run_python: dict`: A dictionary with the following keys: * `uri: str`: The URI of the Python file to run. * `args: list[str]`: A list of arguments to pass to the Python file. diff --git a/airflow_providers_wherobots/operators/run.py b/airflow_providers_wherobots/operators/run.py index 7c12545..b2ed722 100644 --- a/airflow_providers_wherobots/operators/run.py +++ b/airflow_providers_wherobots/operators/run.py @@ -42,6 +42,7 @@ def __init__( polling_interval: int = 20, wherobots_conn_id: str = DEFAULT_CONN_ID, poll_logs: bool = False, + timeout_seconds: int = 3600, xcom_push: bool = True, **kwargs, ): @@ -50,6 +51,7 @@ def __init__( self.run_payload: dict[str, Any] = { "runtime": runtime.value, "name": name or self.default_run_name, + "timeoutSeconds": timeout_seconds, } if run_python: self.run_payload["runPython"] = run_python diff --git a/tests/unit_tests/hooks/test_rest_api.py b/tests/unit_tests/hooks/test_rest_api.py index 4237445..5a2bee6 100644 --- a/tests/unit_tests/hooks/test_rest_api.py +++ b/tests/unit_tests/hooks/test_rest_api.py @@ -115,6 +115,7 @@ def test_create_run(self, test_default_conn) -> None: "args": ["arg1", "arg2"], "entrypoint": "src.main", }, + "timeoutSeconds": 5000, } responses.add( responses.POST,