fix: align registered Runtime & appended Worker timeout (#653)

* fix: align registered Runtime & appended Worker timeout

Signed-off-by: Keming <[email protected]>

* fix timeout condition

Signed-off-by: Keming <[email protected]>

* add fish to envd

Signed-off-by: Keming <[email protected]>

---------

Signed-off-by: Keming <[email protected]>

Author: kemingy

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "mosec"
-version = "0.9.4"
+version = "0.9.5"
 authors = ["Keming <[email protected]>", "Zichen <[email protected]>"]
 edition = "2024"
 license = "Apache-2.0"
@@ -49,8 +49,8 @@ serde = "1.0"
 serde_json = "1.0"
 utoipa = "5.3"
 utoipa-swagger-ui = { version = "9", features = ["axum"] }
-tower = "0.5.1"
-tower-http = { version = "0.6.4", features = [
+tower = "0.5.2"
+tower-http = { version = "0.6.6", features = [
   "compression-zstd",
   "decompression-zstd",
   "compression-gzip",

Cargo.toml

@@ -29,8 +29,8 @@ def jax():
 
 def build():
     base(dev=True, image="ubuntu:22.04")
-    install.conda()
-    install.python()
+    install.uv()
+    shell("fish")
     rust()
     runtime.init(["make install"])
 

build.envd

@@ -50,11 +50,11 @@ class State(enum.IntFlag):
 
 
 @contextmanager
-def set_mosec_timeout(duration: int):
+def set_mosec_timeout(duration: float):
     """Context manager to set a timeout for a code block.
 
     Args:
-        duration (float): the duration in seconds before timing out
+        duration: the duration in seconds before timing out
 
     """
 
@@ -64,11 +64,11 @@ def handler(signum, frame):
         )
 
     signal.signal(signal.SIGALRM, handler)
-    signal.alarm(duration)
+    signal.setitimer(signal.ITIMER_REAL, duration)
     try:
         yield
     finally:
-        signal.alarm(0)
+        signal.setitimer(signal.ITIMER_REAL, 0)
 
 
 class FakeSemaphore:
@@ -108,7 +108,7 @@ def __init__(
         socket_prefix: str,
         stage_name: str,
         worker_id: int,
-        timeout: int,
+        timeout: float,
     ):
         """Initialize the mosec coordinator.
 

mosec/coordinator.py

@@ -85,6 +85,14 @@ def validate_int_ge(number, name, threshold=1):
     assert number >= threshold, f"{name} must be no less than {threshold}"
 
 
+def validate_float_ge(number, name, threshold=0.0):
+    """Validate float number is greater than threshold."""
+    assert isinstance(number, float), (
+        f"{name} must be float but you give {type(number)}"
+    )
+    assert number >= threshold, f"{name} must be no less than {threshold}"
+
+
 def validate_str_dict(dictionary: Dict):
     """Validate keys and values of the dictionary is string type."""
     for key, value in dictionary.items():

mosec/env.py

@@ -24,7 +24,7 @@
 from typing import Callable, Dict, Iterable, List, Optional, Type, Union, cast
 
 from mosec.coordinator import Coordinator
-from mosec.env import env_var_context, validate_env, validate_int_ge
+from mosec.env import env_var_context, validate_env, validate_float_ge, validate_int_ge
 from mosec.log import get_internal_logger
 from mosec.utils import get_mosec_path
 from mosec.worker import Worker
@@ -48,23 +48,24 @@ def __init__(
         worker: Type[Worker],
         num: int = 1,
         max_batch_size: int = 1,
-        max_wait_time: int = 10,
-        timeout: int = 3,
+        max_wait_time: int = 0,
+        timeout: float = 0.0,
         start_method: str = "spawn",
         env: Union[None, List[Dict[str, str]]] = None,
     ):
         """Initialize the mosec coordinator.
 
         Args:
-            worker (Worker): subclass of `mosec.Worker` implemented by users.
-            num (int): number of workers
+            worker: subclass of `mosec.Worker` implemented by users.
+            num: number of workers
             max_batch_size: the maximum batch size allowed (>=1), will enable the
                 dynamic batching if it > 1
-            max_wait_time (int): the maximum wait time (millisecond)
+            max_wait_time: the maximum wait time (millisecond)
                 for dynamic batching, needs to be used with `max_batch_size`
                 to enable the feature. If not configure, will use the CLI
                 argument `--wait` (default=10ms)
-            timeout (int): timeout (second) for the `forward` function.
+            timeout: timeout (second) for the `forward` function.
+                If not set, will use the CLI argument `--timeout` (default=3s)
             start_method: the process starting method ("spawn" or "fork")
             env: the environment variables to set before starting the process
 
@@ -73,7 +74,7 @@ def __init__(
         self.num = num
         self.max_batch_size = max_batch_size
         self.max_wait_time = max_wait_time
-        self.timeout = timeout
+        self.timeout_sec = float(timeout)
         self.start_method = start_method
         self.env = env
 
@@ -84,6 +85,13 @@ def __init__(
 
         self._validate()
 
+    def _align_timeout_wait(self, timeout_sec: float, max_wait_time: int):
+        """Align the ``timeout`` and ``max_wait_time`` arguments."""
+        self.timeout_sec = self.timeout_sec if self.timeout_sec > 0 else timeout_sec
+        self.max_wait_time = (
+            self.max_wait_time if self.max_wait_time > 0 else max_wait_time
+        )
+
     @staticmethod
     def _process_healthy(process: Union[BaseProcess, None]) -> bool:
         """Check if the child process is healthy.
@@ -118,7 +126,7 @@ def _start_process(
                 work_path,
                 self.name,
                 worker_id + 1,
-                self.timeout,
+                self.timeout_sec,
             ),
             daemon=True,
         )
@@ -173,7 +181,7 @@ def _validate(self):
         validate_int_ge(self.num, "worker number")
         validate_int_ge(self.max_batch_size, "maximum batch size")
         validate_int_ge(self.max_wait_time, "maximum wait time", 0)
-        validate_int_ge(self.timeout, "forward timeout", 0)
+        validate_float_ge(self.timeout_sec, "forward timeout", 0)
         assert self.start_method in NEW_PROCESS_METHOD, (
             f"start method must be one of {NEW_PROCESS_METHOD}"
         )