Wake the local On Premise worker after the CLI creates its job

cli.py

@@ -13,13 +13,15 @@
 import typing
 import zipfile
 from collections.abc import Sequence
-from contextlib import ExitStack
+from contextlib import ExitStack, suppress
 from importlib.metadata import PackageNotFoundError, version
 from pathlib import Path
 from tarfile import is_tarfile
 from tempfile import TemporaryDirectory
 from typing import Any, BinaryIO, Literal, get_args, get_origin, get_type_hints
 
+import httpx
+
 from . import NOT_GIVEN, APIConnectionError, APIError, NotGiven, Platform
 from ._cli_metadata import MULTIPART_FILES
 
@@ -33,6 +35,23 @@ def platform_url() -> str:
     return os.getenv("ULTRALYTICS_PLATFORM_URL", "https://platform.ultralytics.com").rstrip("/")
 
 
+def wake_worker(response: httpx.Response) -> None:
+    """Have an On Premise worker on this machine claim a just-created ingest or training job instead of polling.
+
+    Responses do not say where a job runs, so every one wakes; without a worker, or an old one, nothing changes.
+    """
+    request = response.request
+    if response.is_success and request.method == "POST" and request.url.path.endswith(("/ingest", "/training/start")):
+        # stream returns after the headers, so a listener that drips a body can't hold the CLI open
+        with (
+            suppress(httpx.HTTPError),
+            httpx.stream(
+                "GET", "http://127.0.0.1:8765/wake", headers={"Origin": platform_url()}, timeout=0.5, trust_env=False
+            ),
+        ):
+            pass
+
+
 def kinds(annotation: Any) -> tuple[set[str], list | None]:
     """Map a generated SDK annotation to JSON kinds and literal choices; empty kinds accept any JSON value."""
     if annotation is Any:
@@ -715,7 +734,8 @@ def dispatch(tokens: list[str]) -> int:
     """Execute one SDK operation, inferring only a missing resource-path owner."""
     tokens = list(tokens)
     with ExitStack() as stack:
-        client = stack.enter_context(Platform(base_url=platform_url()))
+        http_client = httpx.Client(timeout=60.0, event_hooks={"response": [wake_worker]})
+        client = stack.enter_context(Platform(base_url=platform_url(), http_client=http_client))
         if tokens and tokens[0] in CLOUD_COMMANDS:
             if len(tokens) == 1 or any(token in {"help", "--help", "-h"} for token in tokens):
                 print(inspect.getdoc(CLOUD_COMMANDS[tokens[0]]))

sdk/python/src/ultralytics_platform/cli.py

@@ -13,13 +13,15 @@
 import typing
 import zipfile
 from collections.abc import Sequence
-from contextlib import ExitStack
+from contextlib import ExitStack, suppress
 from importlib.metadata import PackageNotFoundError, version
 from pathlib import Path
 from tarfile import is_tarfile
 from tempfile import TemporaryDirectory
 from typing import Any, BinaryIO, Literal, get_args, get_origin, get_type_hints
 
+import httpx
+
 from . import NOT_GIVEN, APIConnectionError, APIError, NotGiven, Platform
 from ._cli_metadata import MULTIPART_FILES
 
@@ -33,6 +35,23 @@ def platform_url() -> str:
     return os.getenv("ULTRALYTICS_PLATFORM_URL", "https://platform.ultralytics.com").rstrip("/")
 
 
+def wake_worker(response: httpx.Response) -> None:
+    """Have an On Premise worker on this machine claim a just-created ingest or training job instead of polling.
+
+    Responses do not say where a job runs, so every one wakes; without a worker, or an old one, nothing changes.
+    """
+    request = response.request
+    if response.is_success and request.method == "POST" and request.url.path.endswith(("/ingest", "/training/start")):
+        # stream returns after the headers, so a listener that drips a body can't hold the CLI open
+        with (
+            suppress(httpx.HTTPError),
+            httpx.stream(
+                "GET", "http://127.0.0.1:8765/wake", headers={"Origin": platform_url()}, timeout=0.5, trust_env=False
+            ),
+        ):
+            pass
+
+
 def kinds(annotation: Any) -> tuple[set[str], list | None]:
     """Map a generated SDK annotation to JSON kinds and literal choices; empty kinds accept any JSON value."""
     if annotation is Any:
@@ -715,7 +734,8 @@ def dispatch(tokens: list[str]) -> int:
     """Execute one SDK operation, inferring only a missing resource-path owner."""
     tokens = list(tokens)
     with ExitStack() as stack:
-        client = stack.enter_context(Platform(base_url=platform_url()))
+        http_client = httpx.Client(timeout=60.0, event_hooks={"response": [wake_worker]})
+        client = stack.enter_context(Platform(base_url=platform_url(), http_client=http_client))
         if tokens and tokens[0] in CLOUD_COMMANDS:
             if len(tokens) == 1 or any(token in {"help", "--help", "-h"} for token in tokens):
                 print(inspect.getdoc(CLOUD_COMMANDS[tokens[0]]))

tests/test_cli.py

@@ -17,7 +17,7 @@
 
 
 @contextmanager
-def loopback(reply):
+def loopback(reply, port=0):
     """Serve `reply(handler, body) -> (status, JSON payload or raw bytes)` on loopback and yield its URL."""
 
     class Receiver(BaseHTTPRequestHandler):
@@ -35,7 +35,7 @@ def do_GET(self):
 
         do_POST = do_PATCH = do_GET
 
-    server = ThreadingHTTPServer(("127.0.0.1", 0), Receiver)
+    server = ThreadingHTTPServer(("127.0.0.1", port), Receiver)
     thread = Thread(target=server.serve_forever, daemon=True)
     thread.start()
     try:
@@ -135,8 +135,22 @@ def run(*args, input=None):
         count = len(requests)
         result = run("cloud", "models", "update", "project=p", "model=m", "metadata=@-", "train_args=@-", input="{}")
         assert result.returncode == 2 and len(requests) == count
-        failed = run("cloud", "training", "start", "model_id=m", 'train_args={"epochs":100}')
-        assert failed.returncode == 1 and "HTTP 422): model and data are required" in failed.stderr
+        wakes = []
+
+        def wake(handler, body):
+            wakes.append((handler.path, handler.headers["Origin"]))
+            return 204, b""
+
+        with loopback(wake, 8765):  # the On Premise worker's fixed local port
+            failed = run("cloud", "training", "start", "model_id=m", 'train_args={"epochs":100}')
+            assert failed.returncode == 1 and "HTTP 422): model and data are required" in failed.stderr
+            for arguments in (
+                ["training", "start", "model_id=m", "train_args={}"],
+                ["datasets", "ingest", "dataset=d", "body={}"],
+            ):
+                result = run("cloud", *arguments)
+                assert result.returncode == 0 and json.loads(result.stdout) == {"ok": True}
+        assert wakes == [("/wake", url)] * 2  # only successful job creations wake, not failures or the owner lookup
         env["ULTRALYTICS_API_KEY"] = "invalid"
         failed = run("cloud", "account", "summary")
         assert failed.returncode == 1 and "sensitive" not in failed.stderr