From 519e943b09fbf7f5a8626ecf40d89d4634a7af97 Mon Sep 17 00:00:00 2001
From: Gabefire <33893811+Gabefire@users.noreply.github.com>
Date: Tue, 19 Dec 2023 14:47:03 -0600
Subject: [PATCH] modified wait till done method

---
 labelbox/schema/task.py | 17 +++++++++++------
 1 file changed, 11 insertions(+), 6 deletions(-)

diff --git a/labelbox/schema/task.py b/labelbox/schema/task.py
index 717f12ff5..8ba91470c 100644
--- a/labelbox/schema/task.py
+++ b/labelbox/schema/task.py
@@ -58,14 +58,19 @@ def refresh(self) -> None:
         for field in self.fields():
             setattr(self, field.name, getattr(tasks[0], field.name))
 
-    def wait_till_done(self, timeout_seconds: int = 300) -> None:
+    def wait_till_done(self,
+                       timeout_seconds: float = 300.0,
+                       check_frequency: float = 2.0) -> None:
         """ Waits until the task is completed. Periodically queries the server
-        to update the task attributes.
+            to update the task attributes.
 
-        Args:
-            timeout_seconds (float): Maximum time this method can block, in seconds. Defaults to five minutes.
-        """
-        check_frequency = 2  # frequency of checking, in seconds
+            Args:
+                timeout_seconds (float): Maximum time this method can block, in seconds. Defaults to five minutes.
+                check_frequency (float): Frequency of queries to server to update the task attributes, in seconds. Defaults to two seconds. Minimal value is two seconds. 
+            """
+        if check_frequency < 2.0:
+            raise ValueError(
+                "Expected check frequency to be two seconds or more")
         while timeout_seconds > 0:
             if self.status != "IN_PROGRESS":
                 # self.errors fetches the error content.