mmo-runtime 0.1.3__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- mmo_runtime-0.1.3/PKG-INFO +71 -0
- mmo_runtime-0.1.3/README.md +54 -0
- mmo_runtime-0.1.3/pyproject.toml +49 -0
- mmo_runtime-0.1.3/setup.cfg +4 -0
- mmo_runtime-0.1.3/src/mmo_runtime/__init__.py +55 -0
- mmo_runtime-0.1.3/src/mmo_runtime/_version.py +24 -0
- mmo_runtime-0.1.3/src/mmo_runtime/client.py +294 -0
- mmo_runtime-0.1.3/src/mmo_runtime/config.py +90 -0
- mmo_runtime-0.1.3/src/mmo_runtime/context.py +109 -0
- mmo_runtime-0.1.3/src/mmo_runtime/errors.py +125 -0
- mmo_runtime-0.1.3/src/mmo_runtime/events.py +31 -0
- mmo_runtime-0.1.3/src/mmo_runtime/layout.py +46 -0
- mmo_runtime-0.1.3/src/mmo_runtime/logging.py +116 -0
- mmo_runtime-0.1.3/src/mmo_runtime/manifest.py +53 -0
- mmo_runtime-0.1.3/src/mmo_runtime/py.typed +0 -0
- mmo_runtime-0.1.3/src/mmo_runtime/server.py +293 -0
- mmo_runtime-0.1.3/src/mmo_runtime/transport.py +19 -0
- mmo_runtime-0.1.3/src/mmo_runtime/worker.py +254 -0
- mmo_runtime-0.1.3/src/mmo_runtime.egg-info/PKG-INFO +71 -0
- mmo_runtime-0.1.3/src/mmo_runtime.egg-info/SOURCES.txt +28 -0
- mmo_runtime-0.1.3/src/mmo_runtime.egg-info/dependency_links.txt +1 -0
- mmo_runtime-0.1.3/src/mmo_runtime.egg-info/requires.txt +7 -0
- mmo_runtime-0.1.3/src/mmo_runtime.egg-info/top_level.txt +1 -0
- mmo_runtime-0.1.3/tests/test_completion.py +368 -0
- mmo_runtime-0.1.3/tests/test_config.py +107 -0
- mmo_runtime-0.1.3/tests/test_context.py +87 -0
- mmo_runtime-0.1.3/tests/test_errors.py +92 -0
- mmo_runtime-0.1.3/tests/test_layout.py +66 -0
- mmo_runtime-0.1.3/tests/test_manifest.py +77 -0
- mmo_runtime-0.1.3/tests/test_worker.py +223 -0
|
@@ -0,0 +1,71 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: mmo-runtime
|
|
3
|
+
Version: 0.1.3
|
|
4
|
+
Summary: Runtime library for MMO gRPC tools
|
|
5
|
+
Author-email: tool-mmo <tranduythuc2003@gmail.com>
|
|
6
|
+
License: MIT
|
|
7
|
+
Project-URL: Homepage, https://gitlab.com/tool-mmo3/base/toolkit-mmo
|
|
8
|
+
Project-URL: Repository, https://gitlab.com/tool-mmo3/base/toolkit-mmo
|
|
9
|
+
Requires-Python: >=3.11
|
|
10
|
+
Description-Content-Type: text/markdown
|
|
11
|
+
Provides-Extra: dev
|
|
12
|
+
Requires-Dist: pytest>=8.2; extra == "dev"
|
|
13
|
+
Requires-Dist: pytest-cov>=5.0; extra == "dev"
|
|
14
|
+
Requires-Dist: ruff>=0.5; extra == "dev"
|
|
15
|
+
Requires-Dist: pyright>=1.1; extra == "dev"
|
|
16
|
+
Requires-Dist: import-linter>=2.4; extra == "dev"
|
|
17
|
+
|
|
18
|
+
# mmo-runtime
|
|
19
|
+
|
|
20
|
+
Runtime library for MMO gRPC tools.
|
|
21
|
+
|
|
22
|
+
## Worker execution
|
|
23
|
+
|
|
24
|
+
`stream_in_worker(workflow, request, ctx, *, pool=None, timeout=None)` yields
|
|
25
|
+
`ProgressEvent` objects in order, then one `ResultEnvelope` or `ErrorEnvelope`.
|
|
26
|
+
Even without a supplied pool the workflow runs on a worker thread, so progress
|
|
27
|
+
is available before completion. Progress is buffered up to 64 events; slow
|
|
28
|
+
consumers apply backpressure rather than dropping events or growing memory
|
|
29
|
+
without bound. Close abandoned iterators to release blocked producers.
|
|
30
|
+
|
|
31
|
+
Timeout is measured from the first iterator advance and includes queue time.
|
|
32
|
+
`None` disables the deadline; zero (or a negative value) expires immediately.
|
|
33
|
+
Cancellation, timeout, and iterator closure signal `ctx.cancel` and cancel jobs
|
|
34
|
+
which have not started. Running workflows must cooperatively check the token
|
|
35
|
+
and bound their own blocking I/O. Python threads are **not** forcibly terminated;
|
|
36
|
+
an uncooperative job can retain its worker slot and delay process exit.
|
|
37
|
+
|
|
38
|
+
`run_in_worker` retains its legacy `(terminal, progress_list)` API, including
|
|
39
|
+
inline execution without a pool and pooled cancellation/admission exceptions.
|
|
40
|
+
Its accumulated progress list is intentionally not memory-bounded; use the
|
|
41
|
+
streaming API for long-running jobs. `WorkerPool.shutdown(cancel_futures=True)`
|
|
42
|
+
cancels waiting jobs; it does not interrupt running jobs.
|
|
43
|
+
|
|
44
|
+
## Transport and lifecycle
|
|
45
|
+
|
|
46
|
+
`ToolClient.channel` is public for constructing service stubs. Unary and stream
|
|
47
|
+
calls accept `deadline` (including zero) and keyword-only `metadata`, a sequence
|
|
48
|
+
of `(key, value)` pairs. Unknown error-envelope codes normalize to `INTERNAL`.
|
|
49
|
+
Unavailable transports map to retryable `PROVIDER_ERROR`; resource exhaustion,
|
|
50
|
+
invalid requests, cancellation, and expired deadlines retain their error codes.
|
|
51
|
+
Health polling limits each probe and sleep to the remaining polling budget.
|
|
52
|
+
|
|
53
|
+
Plain Unix socket paths, explicit `unix://` targets, TCP `host:port` (including
|
|
54
|
+
IPv6), and resolver targets are normalized without duplicate Unix prefixes.
|
|
55
|
+
The default on Windows is `127.0.0.1:50051`; choose explicit per-tool ports to
|
|
56
|
+
avoid collisions. Unix keeps `/tmp/tool.sock` as its default.
|
|
57
|
+
|
|
58
|
+
The server tracks the full handler lifetime of unary and streaming application
|
|
59
|
+
RPCs. Health probes do not keep an otherwise idle server alive. The idle period
|
|
60
|
+
restarts after the last active handler exits. Shutdown is one-shot, marks
|
|
61
|
+
health not-serving, signals the server cancellation token, waits for gRPC's
|
|
62
|
+
graceful stop, then cleans up transport resources and invokes `on_stop` once.
|
|
63
|
+
Signal handlers are installed only on the Unix main thread and restored on exit.
|
|
64
|
+
An instance is single-use; create a new instance to restart.
|
|
65
|
+
|
|
66
|
+
The server cancellation token is available for embedded workflows; generated
|
|
67
|
+
servicers use their own per-RPC token and gRPC cancellation callback. Graceful
|
|
68
|
+
shutdown therefore lets existing generated RPCs finish until grace expires.
|
|
69
|
+
Neither `job_timeout_sec` nor the server token is automatically injected into
|
|
70
|
+
arbitrary registered handlers: pass explicit worker deadlines and cancellation
|
|
71
|
+
tokens when implementing custom handlers.
|
|
@@ -0,0 +1,54 @@
|
|
|
1
|
+
# mmo-runtime
|
|
2
|
+
|
|
3
|
+
Runtime library for MMO gRPC tools.
|
|
4
|
+
|
|
5
|
+
## Worker execution
|
|
6
|
+
|
|
7
|
+
`stream_in_worker(workflow, request, ctx, *, pool=None, timeout=None)` yields
|
|
8
|
+
`ProgressEvent` objects in order, then one `ResultEnvelope` or `ErrorEnvelope`.
|
|
9
|
+
Even without a supplied pool the workflow runs on a worker thread, so progress
|
|
10
|
+
is available before completion. Progress is buffered up to 64 events; slow
|
|
11
|
+
consumers apply backpressure rather than dropping events or growing memory
|
|
12
|
+
without bound. Close abandoned iterators to release blocked producers.
|
|
13
|
+
|
|
14
|
+
Timeout is measured from the first iterator advance and includes queue time.
|
|
15
|
+
`None` disables the deadline; zero (or a negative value) expires immediately.
|
|
16
|
+
Cancellation, timeout, and iterator closure signal `ctx.cancel` and cancel jobs
|
|
17
|
+
which have not started. Running workflows must cooperatively check the token
|
|
18
|
+
and bound their own blocking I/O. Python threads are **not** forcibly terminated;
|
|
19
|
+
an uncooperative job can retain its worker slot and delay process exit.
|
|
20
|
+
|
|
21
|
+
`run_in_worker` retains its legacy `(terminal, progress_list)` API, including
|
|
22
|
+
inline execution without a pool and pooled cancellation/admission exceptions.
|
|
23
|
+
Its accumulated progress list is intentionally not memory-bounded; use the
|
|
24
|
+
streaming API for long-running jobs. `WorkerPool.shutdown(cancel_futures=True)`
|
|
25
|
+
cancels waiting jobs; it does not interrupt running jobs.
|
|
26
|
+
|
|
27
|
+
## Transport and lifecycle
|
|
28
|
+
|
|
29
|
+
`ToolClient.channel` is public for constructing service stubs. Unary and stream
|
|
30
|
+
calls accept `deadline` (including zero) and keyword-only `metadata`, a sequence
|
|
31
|
+
of `(key, value)` pairs. Unknown error-envelope codes normalize to `INTERNAL`.
|
|
32
|
+
Unavailable transports map to retryable `PROVIDER_ERROR`; resource exhaustion,
|
|
33
|
+
invalid requests, cancellation, and expired deadlines retain their error codes.
|
|
34
|
+
Health polling limits each probe and sleep to the remaining polling budget.
|
|
35
|
+
|
|
36
|
+
Plain Unix socket paths, explicit `unix://` targets, TCP `host:port` (including
|
|
37
|
+
IPv6), and resolver targets are normalized without duplicate Unix prefixes.
|
|
38
|
+
The default on Windows is `127.0.0.1:50051`; choose explicit per-tool ports to
|
|
39
|
+
avoid collisions. Unix keeps `/tmp/tool.sock` as its default.
|
|
40
|
+
|
|
41
|
+
The server tracks the full handler lifetime of unary and streaming application
|
|
42
|
+
RPCs. Health probes do not keep an otherwise idle server alive. The idle period
|
|
43
|
+
restarts after the last active handler exits. Shutdown is one-shot, marks
|
|
44
|
+
health not-serving, signals the server cancellation token, waits for gRPC's
|
|
45
|
+
graceful stop, then cleans up transport resources and invokes `on_stop` once.
|
|
46
|
+
Signal handlers are installed only on the Unix main thread and restored on exit.
|
|
47
|
+
An instance is single-use; create a new instance to restart.
|
|
48
|
+
|
|
49
|
+
The server cancellation token is available for embedded workflows; generated
|
|
50
|
+
servicers use their own per-RPC token and gRPC cancellation callback. Graceful
|
|
51
|
+
shutdown therefore lets existing generated RPCs finish until grace expires.
|
|
52
|
+
Neither `job_timeout_sec` nor the server token is automatically injected into
|
|
53
|
+
arbitrary registered handlers: pass explicit worker deadlines and cancellation
|
|
54
|
+
tokens when implementing custom handlers.
|
|
@@ -0,0 +1,49 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "mmo-runtime"
|
|
3
|
+
description = "Runtime library for MMO gRPC tools"
|
|
4
|
+
readme = "README.md"
|
|
5
|
+
requires-python = ">=3.11"
|
|
6
|
+
license = {text = "MIT"}
|
|
7
|
+
authors = [
|
|
8
|
+
{name = "tool-mmo", email = "tranduythuc2003@gmail.com"},
|
|
9
|
+
]
|
|
10
|
+
dynamic = ["version"]
|
|
11
|
+
|
|
12
|
+
[project.urls]
|
|
13
|
+
Homepage = "https://gitlab.com/tool-mmo3/base/toolkit-mmo"
|
|
14
|
+
Repository = "https://gitlab.com/tool-mmo3/base/toolkit-mmo"
|
|
15
|
+
|
|
16
|
+
[project.optional-dependencies]
|
|
17
|
+
dev = [
|
|
18
|
+
"pytest>=8.2",
|
|
19
|
+
"pytest-cov>=5.0",
|
|
20
|
+
"ruff>=0.5",
|
|
21
|
+
"pyright>=1.1",
|
|
22
|
+
"import-linter>=2.4",
|
|
23
|
+
]
|
|
24
|
+
|
|
25
|
+
[build-system]
|
|
26
|
+
requires = ["setuptools>=61.0", "wheel", "setuptools-scm>=8"]
|
|
27
|
+
build-backend = "setuptools.build_meta"
|
|
28
|
+
|
|
29
|
+
[tool.setuptools.packages.find]
|
|
30
|
+
where = ["src"]
|
|
31
|
+
include = ["mmo_runtime*"]
|
|
32
|
+
|
|
33
|
+
[tool.setuptools_scm]
|
|
34
|
+
write_to = "src/mmo_runtime/_version.py"
|
|
35
|
+
|
|
36
|
+
[tool.ruff]
|
|
37
|
+
target-version = "py311"
|
|
38
|
+
line-length = 100
|
|
39
|
+
|
|
40
|
+
[tool.pyright]
|
|
41
|
+
typeCheckingMode = "strict"
|
|
42
|
+
pythonVersion = "3.11"
|
|
43
|
+
|
|
44
|
+
[tool.importlinter]
|
|
45
|
+
root_package = "mmo_runtime"
|
|
46
|
+
include_external_packages = true
|
|
47
|
+
contracts = [
|
|
48
|
+
{name = "no_ai", type = "forbidden", source_modules = ["mmo_runtime"], forbidden_modules = ["torch", "transformers", "onnx", "tensorflow", "numpy", "cv2", "PIL"]},
|
|
49
|
+
]
|
|
@@ -0,0 +1,55 @@
|
|
|
1
|
+
"""mmo-runtime: Runtime library for MMO gRPC tools."""
|
|
2
|
+
|
|
3
|
+
from .client import ToolClient, consume_stream, ensure_running
|
|
4
|
+
from .config import BaseServerConfig, load_config
|
|
5
|
+
from .context import CancelToken, ProgressReporter, RunContext
|
|
6
|
+
from .errors import (
|
|
7
|
+
Cancelled,
|
|
8
|
+
DeadlineExceeded,
|
|
9
|
+
Internal,
|
|
10
|
+
InvalidRequest,
|
|
11
|
+
ProviderError,
|
|
12
|
+
ResourceExhausted,
|
|
13
|
+
ToolError,
|
|
14
|
+
to_grpc_status,
|
|
15
|
+
wrap_unknown,
|
|
16
|
+
)
|
|
17
|
+
from .events import ErrorEnvelope, ProgressEvent, ResultEnvelope
|
|
18
|
+
from .layout import RunLayout
|
|
19
|
+
from .logging import setup_logging
|
|
20
|
+
from .manifest import Manifest
|
|
21
|
+
from .server import serve
|
|
22
|
+
from .worker import run_in_worker
|
|
23
|
+
|
|
24
|
+
try:
|
|
25
|
+
from ._version import __version__
|
|
26
|
+
except ImportError: # pragma: no cover
|
|
27
|
+
__version__ = "0.0.0+unknown"
|
|
28
|
+
|
|
29
|
+
__all__ = [
|
|
30
|
+
"BaseServerConfig",
|
|
31
|
+
"CancelToken",
|
|
32
|
+
"Cancelled",
|
|
33
|
+
"DeadlineExceeded",
|
|
34
|
+
"ErrorEnvelope",
|
|
35
|
+
"Internal",
|
|
36
|
+
"InvalidRequest",
|
|
37
|
+
"Manifest",
|
|
38
|
+
"ProgressEvent",
|
|
39
|
+
"ProgressReporter",
|
|
40
|
+
"ProviderError",
|
|
41
|
+
"ResourceExhausted",
|
|
42
|
+
"ResultEnvelope",
|
|
43
|
+
"RunContext",
|
|
44
|
+
"RunLayout",
|
|
45
|
+
"ToolClient",
|
|
46
|
+
"ToolError",
|
|
47
|
+
"consume_stream",
|
|
48
|
+
"ensure_running",
|
|
49
|
+
"load_config",
|
|
50
|
+
"run_in_worker",
|
|
51
|
+
"serve",
|
|
52
|
+
"setup_logging",
|
|
53
|
+
"to_grpc_status",
|
|
54
|
+
"wrap_unknown",
|
|
55
|
+
]
|
|
@@ -0,0 +1,24 @@
|
|
|
1
|
+
# file generated by vcs-versioning
|
|
2
|
+
# don't change, don't track in version control
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
__all__ = [
|
|
6
|
+
"__version__",
|
|
7
|
+
"__version_tuple__",
|
|
8
|
+
"version",
|
|
9
|
+
"version_tuple",
|
|
10
|
+
"__commit_id__",
|
|
11
|
+
"commit_id",
|
|
12
|
+
]
|
|
13
|
+
|
|
14
|
+
version: str
|
|
15
|
+
__version__: str
|
|
16
|
+
__version_tuple__: tuple[int | str, ...]
|
|
17
|
+
version_tuple: tuple[int | str, ...]
|
|
18
|
+
commit_id: str | None
|
|
19
|
+
__commit_id__: str | None
|
|
20
|
+
|
|
21
|
+
__version__ = version = '0.1.3'
|
|
22
|
+
__version_tuple__ = version_tuple = (0, 1, 3)
|
|
23
|
+
|
|
24
|
+
__commit_id__ = commit_id = None
|
|
@@ -0,0 +1,294 @@
|
|
|
1
|
+
"""gRPC client with deadline, stream consumption, and ensure_running."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import logging
|
|
6
|
+
import os
|
|
7
|
+
import subprocess
|
|
8
|
+
import sys
|
|
9
|
+
import threading
|
|
10
|
+
import time
|
|
11
|
+
from collections.abc import Callable, Iterator, Sequence
|
|
12
|
+
from pathlib import Path
|
|
13
|
+
from typing import TYPE_CHECKING, Any, NoReturn, Self, TypeVar
|
|
14
|
+
|
|
15
|
+
import grpc
|
|
16
|
+
from grpc_health.v1 import health_pb2, health_pb2_grpc
|
|
17
|
+
|
|
18
|
+
from .errors import Cancelled, DeadlineExceeded, ErrorCode, Internal, ToolError
|
|
19
|
+
from .events import ProgressEvent, ResultEnvelope
|
|
20
|
+
from .transport import normalize_target
|
|
21
|
+
|
|
22
|
+
if TYPE_CHECKING:
|
|
23
|
+
from types import TracebackType
|
|
24
|
+
|
|
25
|
+
T = TypeVar("T")
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
def _get_target_address(socket_path: str) -> str:
|
|
29
|
+
"""Determine gRPC target address for the current platform.
|
|
30
|
+
|
|
31
|
+
On Windows ``socket_path`` is a TCP ``host:port`` (grpc Python has no
|
|
32
|
+
``windows://`` named-pipe resolver). On Unix it is a Unix domain socket path.
|
|
33
|
+
"""
|
|
34
|
+
return normalize_target(socket_path)
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
def consume_stream(
|
|
38
|
+
stream: Iterator[Any],
|
|
39
|
+
on_progress: Callable[[ProgressEvent], None] | None = None,
|
|
40
|
+
) -> ResultEnvelope:
|
|
41
|
+
"""
|
|
42
|
+
Consume a server-streaming RPC, yielding progress and returning final result.
|
|
43
|
+
Raises ToolError on error status or error envelope.
|
|
44
|
+
"""
|
|
45
|
+
result: ResultEnvelope = ResultEnvelope(payload=None)
|
|
46
|
+
has_result = False
|
|
47
|
+
for msg in stream:
|
|
48
|
+
# Check for progress/result/error oneof
|
|
49
|
+
if hasattr(msg, "WhichOneof"):
|
|
50
|
+
oneof = msg.WhichOneof("payload")
|
|
51
|
+
if oneof == "progress":
|
|
52
|
+
prog = msg.progress
|
|
53
|
+
event = ProgressEvent(
|
|
54
|
+
step=prog.step,
|
|
55
|
+
done=prog.done,
|
|
56
|
+
total=prog.total if prog.total > 0 else None,
|
|
57
|
+
message=prog.message if prog.message else None,
|
|
58
|
+
meta=dict(prog.meta) if hasattr(prog, "meta") else {},
|
|
59
|
+
)
|
|
60
|
+
if on_progress:
|
|
61
|
+
on_progress(event)
|
|
62
|
+
elif oneof == "result":
|
|
63
|
+
result = ResultEnvelope(
|
|
64
|
+
payload=msg.result.payload if hasattr(msg.result, "payload") else msg.result,
|
|
65
|
+
artifacts=list(msg.result.artifacts)
|
|
66
|
+
if hasattr(msg.result, "artifacts")
|
|
67
|
+
else [],
|
|
68
|
+
meta=dict(msg.result.meta) if hasattr(msg.result, "meta") else {},
|
|
69
|
+
)
|
|
70
|
+
has_result = True
|
|
71
|
+
elif oneof == "error":
|
|
72
|
+
err = msg.error
|
|
73
|
+
raise ToolError(
|
|
74
|
+
code=err.code,
|
|
75
|
+
step=err.step if err.step else None,
|
|
76
|
+
provider=err.provider if err.provider else None,
|
|
77
|
+
message=err.message,
|
|
78
|
+
retryable=err.retryable,
|
|
79
|
+
)
|
|
80
|
+
else:
|
|
81
|
+
# Fallback: assume message is result
|
|
82
|
+
result = ResultEnvelope(payload=msg)
|
|
83
|
+
has_result = True
|
|
84
|
+
|
|
85
|
+
if not has_result:
|
|
86
|
+
raise Internal("Stream ended without result")
|
|
87
|
+
return result
|
|
88
|
+
|
|
89
|
+
|
|
90
|
+
class ToolClient:
|
|
91
|
+
"""Client for tool gRPC service."""
|
|
92
|
+
|
|
93
|
+
def __init__(
|
|
94
|
+
self,
|
|
95
|
+
socket_path: str,
|
|
96
|
+
default_deadline: float = 30.0,
|
|
97
|
+
health_check_timeout: float = 5.0,
|
|
98
|
+
):
|
|
99
|
+
self.socket_path = socket_path
|
|
100
|
+
self.default_deadline = default_deadline
|
|
101
|
+
self.health_check_timeout = health_check_timeout
|
|
102
|
+
self._channel: grpc.Channel | None = None
|
|
103
|
+
self._health_stub: health_pb2_grpc.HealthStub | None = None
|
|
104
|
+
self._log = logging.getLogger(__name__)
|
|
105
|
+
self._channel_lock = threading.RLock()
|
|
106
|
+
|
|
107
|
+
@property
|
|
108
|
+
def channel(self) -> grpc.Channel:
|
|
109
|
+
"""Public lazily-created channel for generated service stubs."""
|
|
110
|
+
return self._ensure_channel()
|
|
111
|
+
|
|
112
|
+
def _ensure_channel(self) -> grpc.Channel:
|
|
113
|
+
with self._channel_lock:
|
|
114
|
+
if self._channel is None:
|
|
115
|
+
target = _get_target_address(self.socket_path)
|
|
116
|
+
self._channel = grpc.insecure_channel(target)
|
|
117
|
+
self._health_stub = health_pb2_grpc.HealthStub(self._channel)
|
|
118
|
+
return self._channel
|
|
119
|
+
|
|
120
|
+
def close(self) -> None:
|
|
121
|
+
with self._channel_lock:
|
|
122
|
+
if self._channel is not None:
|
|
123
|
+
self._channel.close()
|
|
124
|
+
self._channel = None
|
|
125
|
+
self._health_stub = None
|
|
126
|
+
|
|
127
|
+
def __enter__(self) -> Self:
|
|
128
|
+
return self
|
|
129
|
+
|
|
130
|
+
def __exit__(
|
|
131
|
+
self,
|
|
132
|
+
exc_type: type[BaseException] | None,
|
|
133
|
+
exc_val: BaseException | None,
|
|
134
|
+
exc_tb: TracebackType | None,
|
|
135
|
+
) -> None:
|
|
136
|
+
self.close()
|
|
137
|
+
|
|
138
|
+
def health_check(
|
|
139
|
+
self,
|
|
140
|
+
service: str = "",
|
|
141
|
+
*,
|
|
142
|
+
timeout: float | None = None,
|
|
143
|
+
) -> health_pb2.HealthCheckResponse.ServingStatus:
|
|
144
|
+
self._ensure_channel()
|
|
145
|
+
try:
|
|
146
|
+
resp: health_pb2.HealthCheckResponse = self._health_stub.Check( # type: ignore[attr-defined]
|
|
147
|
+
health_pb2.HealthCheckRequest(service=service),
|
|
148
|
+
timeout=self.health_check_timeout if timeout is None else timeout,
|
|
149
|
+
)
|
|
150
|
+
# grpc_health pb2 ships untyped; status enum member is Unknown to pyright
|
|
151
|
+
return resp.status # type: ignore[return-value]
|
|
152
|
+
except grpc.RpcError as e:
|
|
153
|
+
if e.code() in (grpc.StatusCode.UNAVAILABLE, grpc.StatusCode.DEADLINE_EXCEEDED):
|
|
154
|
+
return health_pb2.HealthCheckResponse.ServingStatus.SERVICE_UNKNOWN
|
|
155
|
+
raise
|
|
156
|
+
|
|
157
|
+
def wait_healthy(self, timeout: float = 30.0, interval: float = 0.5) -> bool:
|
|
158
|
+
deadline = time.monotonic() + timeout
|
|
159
|
+
while time.monotonic() < deadline:
|
|
160
|
+
remaining = min(timeout, max(0, deadline - time.monotonic()))
|
|
161
|
+
if (
|
|
162
|
+
self.health_check(timeout=min(self.health_check_timeout, remaining))
|
|
163
|
+
== health_pb2.HealthCheckResponse.ServingStatus.SERVING
|
|
164
|
+
):
|
|
165
|
+
return True
|
|
166
|
+
time.sleep(max(0, min(interval, deadline - time.monotonic())))
|
|
167
|
+
return False
|
|
168
|
+
|
|
169
|
+
def call_unary(
|
|
170
|
+
self,
|
|
171
|
+
method: Callable[..., Any],
|
|
172
|
+
request: Any,
|
|
173
|
+
deadline: float | None = None,
|
|
174
|
+
*,
|
|
175
|
+
metadata: Sequence[tuple[str, str | bytes]] | None = None,
|
|
176
|
+
) -> Any:
|
|
177
|
+
"""Call a unary RPC with deadline."""
|
|
178
|
+
self._ensure_channel()
|
|
179
|
+
try:
|
|
180
|
+
kwargs: dict[str, Any] = {
|
|
181
|
+
"timeout": self.default_deadline if deadline is None else deadline,
|
|
182
|
+
}
|
|
183
|
+
if metadata is not None:
|
|
184
|
+
kwargs["metadata"] = metadata
|
|
185
|
+
return method(request, **kwargs)
|
|
186
|
+
except grpc.RpcError as e:
|
|
187
|
+
self._handle_grpc_error(e)
|
|
188
|
+
|
|
189
|
+
def call_stream(
|
|
190
|
+
self,
|
|
191
|
+
method: Callable[..., Iterator[Any]],
|
|
192
|
+
request: Any,
|
|
193
|
+
on_progress: Callable[[ProgressEvent], None] | None = None,
|
|
194
|
+
deadline: float | None = None,
|
|
195
|
+
*,
|
|
196
|
+
metadata: Sequence[tuple[str, str | bytes]] | None = None,
|
|
197
|
+
) -> ResultEnvelope:
|
|
198
|
+
"""Call a server-streaming RPC with deadline."""
|
|
199
|
+
self._ensure_channel()
|
|
200
|
+
try:
|
|
201
|
+
kwargs: dict[str, Any] = {
|
|
202
|
+
"timeout": self.default_deadline if deadline is None else deadline,
|
|
203
|
+
}
|
|
204
|
+
if metadata is not None:
|
|
205
|
+
kwargs["metadata"] = metadata
|
|
206
|
+
stream = method(request, **kwargs)
|
|
207
|
+
try:
|
|
208
|
+
return consume_stream(stream, on_progress)
|
|
209
|
+
finally:
|
|
210
|
+
cancel = getattr(stream, "cancel", None)
|
|
211
|
+
if callable(cancel):
|
|
212
|
+
cancel()
|
|
213
|
+
except grpc.RpcError as e:
|
|
214
|
+
self._handle_grpc_error(e)
|
|
215
|
+
|
|
216
|
+
def _handle_grpc_error(self, e: grpc.RpcError) -> NoReturn:
|
|
217
|
+
status = e.code()
|
|
218
|
+
message = e.details() or f"RPC error: {status}"
|
|
219
|
+
if status == grpc.StatusCode.DEADLINE_EXCEEDED:
|
|
220
|
+
raise DeadlineExceeded(message) from e
|
|
221
|
+
if status == grpc.StatusCode.CANCELLED:
|
|
222
|
+
raise Cancelled(message) from e
|
|
223
|
+
code = {
|
|
224
|
+
grpc.StatusCode.INVALID_ARGUMENT: ErrorCode.INVALID_REQUEST,
|
|
225
|
+
grpc.StatusCode.RESOURCE_EXHAUSTED: ErrorCode.RESOURCE_EXHAUSTED,
|
|
226
|
+
grpc.StatusCode.UNKNOWN: ErrorCode.PROVIDER_ERROR,
|
|
227
|
+
grpc.StatusCode.UNAVAILABLE: ErrorCode.PROVIDER_ERROR,
|
|
228
|
+
}.get(status, ErrorCode.INTERNAL)
|
|
229
|
+
raise ToolError(
|
|
230
|
+
code=code,
|
|
231
|
+
message=message,
|
|
232
|
+
retryable=status in (grpc.StatusCode.RESOURCE_EXHAUSTED, grpc.StatusCode.UNAVAILABLE),
|
|
233
|
+
) from e
|
|
234
|
+
|
|
235
|
+
|
|
236
|
+
# Process management for ensure_running
|
|
237
|
+
_spawn_locks: dict[str, threading.Lock] = {}
|
|
238
|
+
_spawn_locks_lock = threading.Lock()
|
|
239
|
+
|
|
240
|
+
|
|
241
|
+
def get_spawn_lock(socket_path: str) -> threading.Lock:
|
|
242
|
+
with _spawn_locks_lock:
|
|
243
|
+
if socket_path not in _spawn_locks:
|
|
244
|
+
_spawn_locks[socket_path] = threading.Lock()
|
|
245
|
+
return _spawn_locks[socket_path]
|
|
246
|
+
|
|
247
|
+
|
|
248
|
+
def ensure_running(
|
|
249
|
+
tool_name: str,
|
|
250
|
+
command: list[str],
|
|
251
|
+
socket_path: str,
|
|
252
|
+
timeout: float = 30.0,
|
|
253
|
+
env: dict[str, str] | None = None,
|
|
254
|
+
cwd: Path | None = None,
|
|
255
|
+
) -> ToolClient:
|
|
256
|
+
"""
|
|
257
|
+
Ensure tool server is running. Spawns if not healthy.
|
|
258
|
+
Uses a per-socket lock to avoid duplicate spawns.
|
|
259
|
+
"""
|
|
260
|
+
client = ToolClient(socket_path)
|
|
261
|
+
if client.wait_healthy(timeout=2.0):
|
|
262
|
+
return client
|
|
263
|
+
|
|
264
|
+
lock = get_spawn_lock(socket_path)
|
|
265
|
+
with lock:
|
|
266
|
+
# Double-check after acquiring lock
|
|
267
|
+
if client.wait_healthy(timeout=2.0):
|
|
268
|
+
return client
|
|
269
|
+
|
|
270
|
+
# Spawn server
|
|
271
|
+
proc_env = os.environ.copy()
|
|
272
|
+
if env:
|
|
273
|
+
proc_env.update(env)
|
|
274
|
+
# On Windows, use CREATE_NEW_PROCESS_GROUP for proper termination
|
|
275
|
+
creationflags = 0
|
|
276
|
+
if sys.platform == "win32":
|
|
277
|
+
creationflags = subprocess.CREATE_NEW_PROCESS_GROUP
|
|
278
|
+
proc = subprocess.Popen(
|
|
279
|
+
command,
|
|
280
|
+
env=proc_env,
|
|
281
|
+
cwd=cwd,
|
|
282
|
+
creationflags=creationflags,
|
|
283
|
+
)
|
|
284
|
+
|
|
285
|
+
# Wait for healthy
|
|
286
|
+
if not client.wait_healthy(timeout=timeout):
|
|
287
|
+
proc.terminate()
|
|
288
|
+
try:
|
|
289
|
+
proc.wait(timeout=5)
|
|
290
|
+
except subprocess.TimeoutExpired:
|
|
291
|
+
proc.kill()
|
|
292
|
+
raise RuntimeError(f"Tool {tool_name} failed to start within {timeout}s")
|
|
293
|
+
|
|
294
|
+
return client
|
|
@@ -0,0 +1,90 @@
|
|
|
1
|
+
"""Configuration loading with precedence: defaults < yaml < env < CLI."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import os
|
|
6
|
+
import sys
|
|
7
|
+
from pathlib import Path
|
|
8
|
+
from typing import Any
|
|
9
|
+
|
|
10
|
+
from pydantic import Field, field_validator
|
|
11
|
+
from pydantic_settings import BaseSettings, SettingsConfigDict
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
class BaseServerConfig(BaseSettings):
|
|
15
|
+
"""Base configuration for tool servers."""
|
|
16
|
+
|
|
17
|
+
model_config = SettingsConfigDict(
|
|
18
|
+
env_prefix="TOOL_",
|
|
19
|
+
env_nested_delimiter="__",
|
|
20
|
+
case_sensitive=False,
|
|
21
|
+
extra="ignore",
|
|
22
|
+
)
|
|
23
|
+
|
|
24
|
+
# Socket
|
|
25
|
+
socket_path: str = Field(
|
|
26
|
+
default_factory=lambda: "127.0.0.1:50051" if sys.platform == "win32" else "/tmp/tool.sock",
|
|
27
|
+
description="gRPC host:port or Unix domain socket path",
|
|
28
|
+
)
|
|
29
|
+
# Concurrency
|
|
30
|
+
max_concurrent_jobs: int = Field(default=1, ge=1, description="Max concurrent jobs")
|
|
31
|
+
max_queue_size: int = Field(default=4, ge=0, description="Max queued jobs (0 = no queue)")
|
|
32
|
+
|
|
33
|
+
@property
|
|
34
|
+
def thread_pool_size(self) -> int:
|
|
35
|
+
"""Total thread pool size: workers + queue + buffer for health/gRPC internals."""
|
|
36
|
+
return self.max_concurrent_jobs + self.max_queue_size + 4
|
|
37
|
+
|
|
38
|
+
# Timeouts
|
|
39
|
+
idle_timeout_sec: int = Field(
|
|
40
|
+
default=300, ge=0, description="Idle timeout before shutdown (0 = disabled)"
|
|
41
|
+
)
|
|
42
|
+
job_timeout_sec: int = Field(default=0, ge=0, description="Per-job timeout (0 = no limit)")
|
|
43
|
+
# Device
|
|
44
|
+
device: str = Field(default="cpu", description="Compute device (cpu, cuda, mps, etc.)")
|
|
45
|
+
# Paths
|
|
46
|
+
runs_root: Path = Field(default=Path("runs"), description="Root directory for run outputs")
|
|
47
|
+
# Logging
|
|
48
|
+
log_level: str = Field(default="INFO", description="Log level")
|
|
49
|
+
log_json: bool = Field(default=False, description="Output JSON logs")
|
|
50
|
+
|
|
51
|
+
@field_validator("runs_root", mode="before")
|
|
52
|
+
@classmethod
|
|
53
|
+
def expand_runs_root(cls, v: str | Path) -> Path:
|
|
54
|
+
return Path(v).expanduser().resolve()
|
|
55
|
+
|
|
56
|
+
@field_validator("socket_path", mode="before")
|
|
57
|
+
@classmethod
|
|
58
|
+
def expand_socket_path(cls, v: str) -> str:
|
|
59
|
+
return os.path.expanduser(v)
|
|
60
|
+
|
|
61
|
+
|
|
62
|
+
def load_config(
|
|
63
|
+
config_cls: type[BaseServerConfig],
|
|
64
|
+
yaml_path: Path | None = None,
|
|
65
|
+
cli_overrides: dict[str, Any] | None = None,
|
|
66
|
+
) -> BaseServerConfig:
|
|
67
|
+
"""
|
|
68
|
+
Load config with precedence: defaults < yaml < env < CLI.
|
|
69
|
+
"""
|
|
70
|
+
yaml_data: dict[str, Any] = {}
|
|
71
|
+
if yaml_path and yaml_path.exists():
|
|
72
|
+
import yaml
|
|
73
|
+
|
|
74
|
+
with open(yaml_path) as f:
|
|
75
|
+
yaml_data = yaml.safe_load(f) or {}
|
|
76
|
+
|
|
77
|
+
# pydantic-settings ranks init kwargs above env vars, so yaml and CLI
|
|
78
|
+
# (both passed as init kwargs) would each beat env. To get yaml < env < CLI:
|
|
79
|
+
# drop yaml keys that the environment also provides; keep CLI keys so they
|
|
80
|
+
# win over env. Only flat (non-nested) yaml keys are env-matched.
|
|
81
|
+
cli = {k: v for k, v in (cli_overrides or {}).items() if v is not None}
|
|
82
|
+
merged: dict[str, Any] = {**yaml_data, **cli}
|
|
83
|
+
|
|
84
|
+
prefix = str(config_cls.model_config.get("env_prefix", ""))
|
|
85
|
+
env_keys = {k.upper() for k in os.environ}
|
|
86
|
+
for key in yaml_data:
|
|
87
|
+
if key not in cli and f"{prefix}{key}".upper() in env_keys:
|
|
88
|
+
merged.pop(key)
|
|
89
|
+
|
|
90
|
+
return config_cls(**merged)
|