fixes, tests, cleanup

Author: piotrekno1

examples/cli/cli_wrapper_example.py

@@ -13,7 +13,7 @@
 iterations = 500
 
 for _ in range(iterations):
-  with torch.inference_mode():
-      output = model(batch)
+    with torch.inference_mode():
+        output = model(batch)
 
 # print("output shape:", tuple(output.shape))

examples/feedback_example.py

@@ -48,3 +48,5 @@
         print(
             f"run {i + 1}: confidence={confidence:.3f} → {handle.last_inference_id[:8]}..."
         )
+
+client.close()

examples/gguf_example.py

@@ -33,3 +33,5 @@
     result = llm(prompt, max_tokens=128, temperature=0.7)
     text = result["choices"][0]["text"].strip()
     print(f"Q: {prompt}\nA: {text}\n")
+
+client.close()

examples/gguf_gemma_manual_example.py

@@ -114,3 +114,5 @@
     except Exception as exc:
         handle.track_error(error_code="UNKNOWN", error_message=str(exc)[:200])
         raise
+
+client.close()

examples/keras_example.py

@@ -42,3 +42,5 @@
 for _ in range(3):
     output = model.predict(batch, verbose=0)
     print("output shape:", output.shape)
+
+client.close()

examples/onnx_example.py

@@ -25,3 +25,5 @@
 for _ in range(3):
     outputs = session.run(None, {"pixel_values": batch})
     print("output shape:", outputs[0].shape)
+
+client.close()

examples/pytorch_example.py

@@ -40,3 +40,5 @@ def forward(self, x):
     for _ in range(3):
         output = model(batch)
         print("output shape:", output.shape)
+
+client.close()

examples/timm_example.py

@@ -33,3 +33,5 @@
     for _ in range(3):
         output = model(batch)
         print("output shape:", output.shape)
+
+client.close()

tests/test_cli.py

@@ -9,6 +9,7 @@
 from wildedge.integrations.registry import IntegrationSpec
 from wildedge.runtime import bootstrap
 from wildedge.runtime import runner as runtime_runner
+from wildedge.settings import read_runtime_env
 
 
 def _fake_execle(captured: dict):
@@ -55,7 +56,7 @@ def test_cli_run_execs_command_with_env(monkeypatch):
     assert captured["env"][constants.ENV_STRICT_INTEGRATIONS] == "0"
     assert captured["env"][constants.ENV_PRINT_STARTUP_REPORT] == "0"
     assert captured["env"][constants.ENV_FLUSH_TIMEOUT] == str(
-        constants.DEFAULT_SHUTDOWN_FLUSH_TIMEOUT_SEC
+        constants.DEFAULT_RUNTIME_FLUSH_TIMEOUT_SEC
     )
 
 
@@ -148,7 +149,7 @@ def test_install_runtime_requires_dsn(monkeypatch):
         bootstrap.install_runtime()
 
 
-def test_install_runtime_default_flush_timeout_is_shutdown_budget(monkeypatch):
+def test_install_runtime_default_flush_timeout_is_runtime_budget(monkeypatch):
     class FakeWildEdge:
         SUPPORTED_INTEGRATIONS = {"onnx"}
 
@@ -171,7 +172,7 @@ def close(self):  # type: ignore[no-untyped-def]
 
     context = bootstrap.install_runtime()
     try:
-        assert context.flush_timeout == constants.DEFAULT_SHUTDOWN_FLUSH_TIMEOUT_SEC
+        assert context.flush_timeout == constants.DEFAULT_RUNTIME_FLUSH_TIMEOUT_SEC
     finally:
         context.shutdown()
 
@@ -471,3 +472,42 @@ def test_parse_run_args_only_double_dash_raises():
     """['--'] with nothing after raises ValueError."""
     with pytest.raises(ValueError, match="missing command"):
         cli.parse_run_args(["--"])
+
+
+def test_cli_run_default_flush_timeout_is_nonzero(monkeypatch):
+    """CLI must pass a non-zero flush timeout so reservoir events are sent at shutdown."""
+    captured: dict = {}
+    monkeypatch.setattr(cli.os, "execle", _fake_execle(captured))
+    monkeypatch.setattr(cli.shutil, "which", lambda cmd: f"/usr/bin/{cmd}")
+
+    cli.main(["run", "--", "gunicorn", "myapp.wsgi:app"])
+
+    flush_timeout = float(captured["env"][constants.ENV_FLUSH_TIMEOUT])
+    assert flush_timeout > 0, (
+        "flush timeout must be > 0 so reservoir inference events are flushed at shutdown"
+    )
+    assert flush_timeout == constants.DEFAULT_RUNTIME_FLUSH_TIMEOUT_SEC
+
+
+def test_cli_run_flush_timeout_override(monkeypatch):
+    """--flush-timeout arg is forwarded to child process."""
+    captured: dict = {}
+    monkeypatch.setattr(cli.os, "execle", _fake_execle(captured))
+    monkeypatch.setattr(cli.shutil, "which", lambda cmd: f"/usr/bin/{cmd}")
+
+    cli.main(["run", "--flush-timeout", "10.0", "--", "gunicorn", "myapp.wsgi:app"])
+
+    assert captured["env"][constants.ENV_FLUSH_TIMEOUT] == "10.0"
+
+
+def test_read_runtime_env_default_flush_timeout_is_nonzero():
+    """read_runtime_env must default to a non-zero flush timeout when env var is absent."""
+    env = {constants.ENV_DSN: "https://secret@ingest.wildedge.dev/key"}
+    result = read_runtime_env(all_integrations=[], all_hubs=[], environ=env)
+    assert result.flush_timeout > 0
+    assert result.flush_timeout == constants.DEFAULT_RUNTIME_FLUSH_TIMEOUT_SEC
+
+
+def test_runtime_flush_timeout_constant_is_nonzero():
+    """Guard: DEFAULT_RUNTIME_FLUSH_TIMEOUT_SEC must stay > 0."""
+    assert constants.DEFAULT_RUNTIME_FLUSH_TIMEOUT_SEC > 0

wildedge/cli.py

@@ -49,7 +49,7 @@ def build_parser() -> argparse.ArgumentParser:
     run.add_argument(
         "--flush-timeout",
         type=float,
-        default=constants.DEFAULT_SHUTDOWN_FLUSH_TIMEOUT_SEC,
+        default=constants.DEFAULT_RUNTIME_FLUSH_TIMEOUT_SEC,
         help="Flush timeout (seconds) for shutdown.",
     )
     run.add_argument(