diff --git a/.github/workflows/interpreter-example-test.yml b/.github/workflows/interpreter-example-test.yml new file mode 100644 index 0000000000..9ca67828db --- /dev/null +++ b/.github/workflows/interpreter-example-test.yml @@ -0,0 +1,29 @@ +name: Interpreter Example Test + +on: + pull_request: + paths: + - 'examples/run-code-env-inheritance/**' + - '.github/workflows/interpreter-example-test.yml' + push: + branches: [master] + paths: + - 'examples/run-code-env-inheritance/**' + - '.github/workflows/interpreter-example-test.yml' + workflow_dispatch: + +permissions: + contents: read + +jobs: + test: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v6 + - uses: actions/setup-python@v6 + with: + python-version: '3.11' + - name: Install interpreter test dependencies + run: python -m pip install fastapi==0.119.0 uvicorn==0.38.0 httpx==0.28.1 websockets==15.0.1 + - name: Run startup and environment regression tests + run: PYTHONPATH=examples/run-code-env-inheritance/lightweight-code-interpreter python -m unittest discover -s examples/run-code-env-inheritance/tests -v diff --git a/examples/run-code-env-inheritance/Dockerfile b/examples/run-code-env-inheritance/Dockerfile index ce505b55bb..581e5f0c57 100644 --- a/examples/run-code-env-inheritance/Dockerfile +++ b/examples/run-code-env-inheritance/Dockerfile @@ -4,3 +4,10 @@ ARG SANDBOX_CODE_IMAGE=cube-sandbox-cn.tencentcloudcr.com/cube-sandbox/sandbox-c FROM ${SANDBOX_CODE_IMAGE} COPY lightweight-code-interpreter/server.py /opt/lightweight-code-interpreter/server.py +COPY lightweight-code-interpreter/launch.py /opt/lightweight-code-interpreter/launch.py + +RUN test -x /usr/bin/tini && command -v python && command -v jupyter \ + && python -c "import sys; import uvicorn; import jupyter_server; sys.path.insert(0, '/opt/lightweight-code-interpreter'); import server" + +ENTRYPOINT ["/usr/bin/tini", "--"] +CMD ["python", "/opt/lightweight-code-interpreter/launch.py"] diff --git a/examples/run-code-env-inheritance/README.md b/examples/run-code-env-inheritance/README.md index bc8c049673..0e9ac02ad4 100644 --- a/examples/run-code-env-inheritance/README.md +++ b/examples/run-code-env-inheritance/README.md @@ -2,7 +2,7 @@ [中文](README_zh.md) -This example provides a derived image based on the official `sandbox-code` image. It replaces only the lightweight code interpreter so `run_code` can read sandbox-level environment variables injected through `Sandbox.create(envs=...)`, while preserving temporary per-call environment overrides. +This example provides a derived image based on the official `sandbox-code` image. It replaces the lightweight code interpreter and its startup supervision so `run_code` can read sandbox-level environment variables injected through `Sandbox.create(envs=...)`, while preserving temporary per-call environment overrides. ## Background @@ -25,7 +25,21 @@ Sandbox-level values and per-call overrides are applied in a separate background The next execution waits for pending cleanup so per-call values cannot leak into later executions. This lifecycle follows E2B code-interpreter. Sandbox-level environment variables are loaded only once per Jupyter kernel; this example does not provide a general runtime refresh mechanism. -This derived interpreter requires envd's `/envs` endpoint. If envd is unavailable, the endpoint is unsupported, the request exceeds `ENVD_TIMEOUT` (2 seconds by default), or the payload is invalid, `run_code` fails with HTTP 502 instead of falling back to the stock interpreter behavior. Set `ENVD_TIMEOUT` in the image environment to adjust the fetch timeout. +This derived interpreter requires envd's `/envs` endpoint. Startup waits for the endpoint within the context-creation budget and fails if it cannot become ready. After startup, an unavailable endpoint, a request exceeding `ENVD_TIMEOUT` (2 seconds by default), or an invalid payload causes `run_code` to return HTTP 502. Set `ENVD_TIMEOUT` in the image environment to adjust the fetch timeout. + +Jupyter context creation, WebSocket opening, and working-directory initialization each use `CODE_INTERPRETER_CONTEXT_CREATE_TIMEOUT` (60 seconds by default), rather than the HTTP and WebSocket defaults of 5 and 10 seconds during concurrent cold starts. A timed-out creation request is not replayed because the kernel may already exist. If no session ID was returned, that session cannot be reclaimed through this cleanup path and may live until the sandbox shuts down. If the session ID is known, a failed handshake closes the connection and attempts to delete that session. Startup also verifies the required envd `/envs` endpoint within the context-creation budget. Failed startup closes the HTTP client without becoming ready. User-code execution timeouts are unchanged. + +The derived image replaces the inherited entrypoint with tini so the launcher owns dependency startup and does not start a second envd. It honors `ENVD_BIN`, `ENVD_EXTRA_ARGS`, and `ENVD_LOG_FILE=-`. Jupyter is configured to fail on a port conflict instead of selecting another port. + +The launcher binds and listens on the interpreter port before starting envd, Jupyter, or application initialization, then passes that same socket to Uvicorn. Uvicorn normally binds after its lifespan startup; dynamic ports allocated during initialization could otherwise take the default port `49999` and terminate startup with `address already in use`. Dependency processes do not inherit the listener and are terminated and reaped on startup failure or shutdown. The image keeps `/usr/bin/tini` as PID 1 to reap orphaned grandchildren, while the Python launcher manages its direct children. Reserving the port does not mean the application is ready. + +The launcher starts envd with `-isnotfc`, matching CubeSandbox's base-image entrypoint. Without this flag, envd 0.5.11 enables its Firecracker MMDS-dependent HTTP log exporter; in an environment without that metadata, the exporter can keep buffering logs without a size limit. Non-Firecracker mode avoids that exporter while preserving the existing local stdout/stderr log file. It does not suppress command logs, force garbage collection, or change command/file APIs. + +Readiness probes must request the HTTP `/health` endpoint and require a successful +response. The launcher reserves the TCP listener before starting dependencies; +a successful TCP connection alone does not mean startup has completed. +`JUPYTER_LOG_FILE` selects the Jupyter log path (default `/var/log/jupyter.log`); +its parent directory is created before launching the process. ## Build and test the image diff --git a/examples/run-code-env-inheritance/lightweight-code-interpreter/launch.py b/examples/run-code-env-inheritance/lightweight-code-interpreter/launch.py new file mode 100644 index 0000000000..984bb2c41e --- /dev/null +++ b/examples/run-code-env-inheritance/lightweight-code-interpreter/launch.py @@ -0,0 +1,107 @@ +"""Reserve the interpreter port before starting services or Jupyter kernels.""" +import os +from contextlib import ExitStack +from pathlib import Path +import signal +import shlex +import socket +import subprocess +import sys + +import uvicorn + + +def exit_on_signal(signum, frame): + # Uvicorn restores and re-raises captured signals. Exit through Python so + # the launcher can reap children even when it is the container's PID 1. + raise SystemExit(128 + signum) + + +def stop_child(process): + if process.poll() is not None: + return + process.terminate() + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=5) + + +def stop_children(children): + pending_error = sys.exc_info()[0] is not None + first_error = None + for child in reversed(children): + try: + stop_child(child) + except (OSError, subprocess.TimeoutExpired) as error: + print(f"[lightweight-code-interpreter] child cleanup failed: {error}", file=sys.stderr, flush=True) + if first_error is None: + first_error = error + if first_error is not None and not pending_error: + raise first_error + + +def main(): + host = os.getenv("CODE_INTERPRETER_HOST", "0.0.0.0") + port = int(os.getenv("CODE_INTERPRETER_PORT", "49999")) + workdir = Path(os.getenv("CODE_INTERPRETER_WORKDIR", "/workspace")) + jupyter_host = os.getenv("JUPYTER_HOST", "127.0.0.1") + jupyter_port = os.getenv("JUPYTER_PORT", "8888") + envd_port = (os.getenv("ENVD_PORT") or "49983") + url_host = f"[{jupyter_host}]" if ":" in jupyter_host else jupyter_host + os.environ["JUPYTER_BASE_URL"] = f"http://{url_host}:{jupyter_port}" + workdir.mkdir(parents=True, exist_ok=True) + config = uvicorn.Config( + "server:app", host=host, port=port, log_level="info", timeout_keep_alive=640 + ) + family = socket.AF_INET6 if ":" in host else socket.AF_INET + children = [] + # Uvicorn's usual single-worker path binds after the application lifespan. + # Its lifespan creates a Jupyter kernel, whose dynamic ports (or an HTTP + # client's ephemeral port) may otherwise take the still-unbound service port. + with socket.socket(family, socket.SOCK_STREAM) as listener: + listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + listener.bind((host, port)) + listener.listen(socket.SOMAXCONN) + listener.set_inheritable(False) + with ExitStack() as stack: + envd_name = os.getenv("ENVD_LOG_FILE") or "/var/log/envd.log" + jupyter_name = os.getenv("JUPYTER_LOG_FILE") or "/var/log/jupyter.log" + for name in (envd_name, jupyter_name): + if name != "-": + Path(name).parent.mkdir(parents=True, exist_ok=True) + envd_log = None if envd_name == "-" else stack.enter_context(open(envd_name, "ab")) + jupyter_log = None if jupyter_name == "-" else stack.enter_context(open(jupyter_name, "ab")) + previous_term = signal.signal(signal.SIGTERM, exit_on_signal) + try: + children.append(subprocess.Popen( + [(os.getenv("ENVD_BIN") or "/usr/bin/envd"), "-isnotfc", "-port", envd_port] + + shlex.split(os.getenv("ENVD_EXTRA_ARGS", "")), + stdin=subprocess.DEVNULL, stdout=envd_log, + stderr=subprocess.STDOUT, close_fds=True, + )) + print(f"[lightweight-code-interpreter] envd launched on :{envd_port}", flush=True) + children.append(subprocess.Popen( + ["jupyter", "server", "--IdentityProvider.token=", + f"--ServerApp.ip={jupyter_host}", f"--ServerApp.port={jupyter_port}", + "--ServerApp.port_retries=0", + f"--ServerApp.root_dir={workdir}", "--ServerApp.allow_root=True", + "--ServerApp.disable_check_xsrf=True", "--ServerApp.open_browser=False"], + stdin=subprocess.DEVNULL, stdout=jupyter_log, + stderr=subprocess.STDOUT, close_fds=True, + )) + print(f"[lightweight-code-interpreter] Jupyter Server started on {os.environ['JUPYTER_BASE_URL']}", flush=True) + server = uvicorn.Server(config) + server.run(sockets=[listener]) + if not server.started: + raise SystemExit(3) + finally: + try: + stop_children(children) + finally: + signal.signal(signal.SIGTERM, previous_term) + + +if __name__ == "__main__": + main() diff --git a/examples/run-code-env-inheritance/lightweight-code-interpreter/server.py b/examples/run-code-env-inheritance/lightweight-code-interpreter/server.py index 42eb7a28fa..ceb02d18a2 100644 --- a/examples/run-code-env-inheritance/lightweight-code-interpreter/server.py +++ b/examples/run-code-env-inheritance/lightweight-code-interpreter/server.py @@ -21,10 +21,11 @@ JUPYTER_BASE_URL = os.getenv("JUPYTER_BASE_URL", "http://127.0.0.1:8888") JUPYTER_WS_URL = JUPYTER_BASE_URL.replace("http://", "ws://").replace("https://", "wss://") -ENVD_ENVS_URL = f"http://127.0.0.1:{os.getenv('ENVD_PORT', '49983')}/envs" +ENVD_ENVS_URL = f"http://127.0.0.1:{os.getenv('ENVD_PORT') or '49983'}/envs" ENVD_TIMEOUT = float(os.getenv("ENVD_TIMEOUT", "2")) WORKDIR = Path(os.getenv("CODE_INTERPRETER_WORKDIR", "/workspace")).resolve() DEFAULT_LANGUAGE = "python" +CONTEXT_CREATE_TIMEOUT = float(os.getenv("CODE_INTERPRETER_CONTEXT_CREATE_TIMEOUT", "60")) PING_TIMEOUT = 30 logging.basicConfig(level=os.getenv("CODE_INTERPRETER_LOG_LEVEL", "INFO")) @@ -156,13 +157,14 @@ async def connect(self) -> None: ws_logger.setLevel(logging.ERROR) self.ws = await connect( self.ws_url, + open_timeout=CONTEXT_CREATE_TIMEOUT, ping_timeout=PING_TIMEOUT, max_size=None, max_queue=None, logger=ws_logger, ) self.receive_task = asyncio.create_task(self._receive_messages()) - await self._run_background(f"%cd {self.cwd}") + await asyncio.wait_for(self._run_background(f"%cd {self.cwd}"), timeout=CONTEXT_CREATE_TIMEOUT) async def _run_background(self, code: str) -> None: if self.ws is None: @@ -439,6 +441,29 @@ async def _wait_for_jupyter() -> None: raise RuntimeError("Jupyter Server did not become ready") +async def _wait_for_envd(http_client: httpx.AsyncClient) -> None: + async def ready(): + while True: + try: + await _fetch_sandbox_envs(http_client) + return + except HTTPException as exc: + cause = exc.__cause__ + retryable = isinstance(cause, httpx.TransportError) or ( + isinstance(cause, httpx.HTTPStatusError) + and (cause.response.status_code >= 500 or cause.response.status_code == 429) + ) + if not retryable: + reason = f"{exc.detail}: {cause}" if cause else exc.detail + raise RuntimeError(f"envd startup failed: {reason}") from exc + await asyncio.sleep(0.1) + + try: + await asyncio.wait_for(ready(), timeout=CONTEXT_CREATE_TIMEOUT) + except asyncio.TimeoutError as exc: + raise RuntimeError("envd did not become ready") from exc + + contexts: dict[str, KernelContext] = {} default_context_id: str | None = None client: httpx.AsyncClient | None = None @@ -451,6 +476,7 @@ async def create_context(language: str | None = None, cwd: str | None = None) -> response = await client.post( f"{JUPYTER_BASE_URL}/api/sessions", + timeout=CONTEXT_CREATE_TIMEOUT, json={ "path": str(uuid.uuid4()), "kernel": {"name": "python3"}, @@ -464,7 +490,20 @@ async def create_context(language: str | None = None, cwd: str | None = None) -> session = response.json() context_id = session["kernel"]["id"] kernel = KernelContext(context_id, session["id"], resolved_cwd) - await kernel.connect() + try: + await kernel.connect() + except BaseException: + # The session ID is known here, so cleanup is safe without replaying + # creation. Keep the original handshake/initialization exception. + try: + await kernel.close() + except BaseException as cleanup_error: + logger.warning("failed context connection cleanup: %s", cleanup_error) + try: + await client.delete(f"{JUPYTER_BASE_URL}/api/sessions/{session['id']}") + except BaseException as cleanup_error: + logger.warning("failed context session cleanup: %s", cleanup_error) + raise contexts[context_id] = kernel return Context(id=context_id, language=DEFAULT_LANGUAGE, cwd=resolved_cwd) @@ -474,14 +513,18 @@ async def lifespan(app: FastAPI): global client, default_context_id await _wait_for_jupyter() client = httpx.AsyncClient() - default = await create_context(DEFAULT_LANGUAGE, str(WORKDIR)) - default_context_id = default.id try: + await _wait_for_envd(client) + # Kernel startup competes with other guests during a cold-start burst. + # Do not replay POST /sessions on timeout: its result can be unknown. + default = await create_context(DEFAULT_LANGUAGE, str(WORKDIR)) + default_context_id = default.id yield finally: for context in list(contexts.values()): await context.close() contexts.clear() + default_context_id = None if client is not None: await client.aclose() diff --git a/examples/run-code-env-inheritance/tests/test_launch.py b/examples/run-code-env-inheritance/tests/test_launch.py new file mode 100644 index 0000000000..4e21d8a24c --- /dev/null +++ b/examples/run-code-env-inheritance/tests/test_launch.py @@ -0,0 +1,238 @@ +import asyncio +import errno +import os +from pathlib import Path +import signal +import socket +import subprocess +import tempfile +import threading +import unittest +from unittest.mock import Mock, patch + +import launch + + +class InterpreterLaunchTest(unittest.TestCase): + def setUp(self): + self.temp = tempfile.TemporaryDirectory() + self.addCleanup(self.temp.cleanup) + root = Path(self.temp.name) + self.env = patch.dict(os.environ, { + "CODE_INTERPRETER_HOST": "127.0.0.1", "CODE_INTERPRETER_PORT": "0", + "CODE_INTERPRETER_WORKDIR": str(root / "work"), + "ENVD_LOG_FILE": str(root / "envd.log"), + "JUPYTER_LOG_FILE": str(root / "jupyter.log"), + }) + self.env.start() + self.addCleanup(self.env.stop) + + def children(self): + return [Mock(poll=Mock(return_value=None)), Mock(poll=Mock(return_value=None))] + + def test_creates_custom_log_parent_directories(self): + envd = Path(self.temp.name) / 'envd' / 'service.log' + jupyter = Path(self.temp.name) / 'jupyter' / 'service.log' + with patch.dict(os.environ, {'ENVD_LOG_FILE': str(envd), 'JUPYTER_LOG_FILE': str(jupyter)}), \ + patch.object(launch.subprocess, 'Popen', side_effect=self.children()), \ + patch.object(launch.uvicorn, 'Server'): + launch.main() + self.assertTrue(envd.is_file()) + self.assertTrue(jupyter.is_file()) + + def test_empty_log_paths_use_entrypoint_defaults(self): + with patch.dict(os.environ, {'ENVD_LOG_FILE': '', 'JUPYTER_LOG_FILE': ''}), \ + patch.object(launch.Path, 'mkdir'), \ + patch('builtins.open', unittest.mock.mock_open()) as opened, \ + patch.object(launch.subprocess, 'Popen', side_effect=self.children()), \ + patch.object(launch.uvicorn, 'Server'): + launch.main() + self.assertEqual([call.args[0] for call in opened.call_args_list], + ['/var/log/envd.log', '/var/log/jupyter.log']) + + def assert_listening(self, listener): + with socket.create_connection(listener.getsockname(), timeout=1): + accepted, _ = listener.accept() + accepted.close() + + def test_envd_uses_non_firecracker_mode_with_custom_port(self): + children = self.children() + with patch.dict(os.environ, {'ENVD_PORT': '50123'}), \ + patch.object(launch.subprocess, 'Popen', side_effect=children) as start, \ + patch.object(launch.uvicorn, 'Server'): + launch.main() + self.assertEqual(start.call_args_list[0].args[0], + ['/usr/bin/envd', '-isnotfc', '-port', '50123']) + + def test_empty_envd_settings_use_entrypoint_defaults(self): + with patch.dict(os.environ, {'ENVD_BIN': '', 'ENVD_PORT': ''}), \ + patch.object(launch.subprocess, 'Popen', side_effect=self.children()) as start, \ + patch.object(launch.uvicorn, 'Server'): + launch.main() + self.assertEqual(start.call_args_list[0].args[0], + ['/usr/bin/envd', '-isnotfc', '-port', '49983']) + + def test_listener_reserved_before_dependencies_and_lifespan(self): + children = self.children() + listener = None + + def construct_server(config): + server = Mock() + def run(*, sockets): + nonlocal listener + listener = sockets[0] + self.assert_listening(listener) + self.assertFalse(listener.get_inheritable()) + with socket.socket() as contender: + with self.assertRaises(OSError) as caught: + contender.bind(listener.getsockname()) + self.assertEqual(caught.exception.errno, errno.EADDRINUSE) + server.run.side_effect = run + return server + + real_socket = socket.socket + sockets = [] + def new_socket(*args, **kwargs): + result = real_socket(*args, **kwargs);sockets.append(result);return result + def start_child(*args, **kwargs): + self.assertTrue(sockets) + self.assert_listening(sockets[0]) + self.assertTrue(kwargs['close_fds']) + return children.pop(0) + with patch.object(launch.socket, 'socket', side_effect=new_socket), \ + patch.object(launch.subprocess, 'Popen', side_effect=start_child), \ + patch.object(launch.uvicorn, 'Server', side_effect=construct_server): + launch.main() + self.assertEqual(listener.fileno(), -1) + + def test_server_failure_reaps_both_children(self): + children = self.children() + with patch.object(launch.subprocess, 'Popen', side_effect=children), \ + patch.object(launch.uvicorn, 'Server') as server: + server.return_value.run.side_effect = RuntimeError('startup failed') + with self.assertRaisesRegex(RuntimeError, 'startup failed'): + launch.main() + for child in children: + child.terminate.assert_called_once() + child.wait.assert_called_once_with(timeout=5) + + def test_partial_child_launch_failure_reaps_started_child(self): + child = self.children()[0] + with patch.object(launch.subprocess, 'Popen', side_effect=[child, OSError('jupyter unavailable')]), \ + patch.object(launch.uvicorn, 'Server') as server: + with self.assertRaisesRegex(OSError, 'jupyter unavailable'): + launch.main() + child.terminate.assert_called_once() + server.assert_not_called() + + def test_failed_lifespan_returns_nonzero_and_reaps_children(self): + children = self.children() + with patch.object(launch.subprocess, 'Popen', side_effect=children), \ + patch.object(launch.uvicorn, 'Server') as server: + server.return_value.started = False + with self.assertRaises(SystemExit) as caught: + launch.main() + self.assertEqual(caught.exception.code, 3) + for child in children: + child.terminate.assert_called_once() + + def test_bind_conflict_does_not_start_children(self): + with socket.socket() as occupied: + occupied.bind(('127.0.0.1', 0)) + occupied.listen() + with patch.dict(os.environ, {'CODE_INTERPRETER_PORT': str(occupied.getsockname()[1])}), \ + patch.object(launch.subprocess, 'Popen') as start: + with self.assertRaises(OSError) as caught: + launch.main() + self.assertEqual(caught.exception.errno, errno.EADDRINUSE) + start.assert_not_called() + + def test_cleanup_failure_still_reaps_other_child(self): + children = self.children() + children[1].terminate.side_effect = OSError('terminate failed') + with self.assertRaisesRegex(OSError, 'terminate failed'): + launch.stop_children(children) + children[0].terminate.assert_called_once() + + def test_cleanup_failure_preserves_startup_failure(self): + children = self.children() + children[1].terminate.side_effect = OSError('terminate failed') + with patch.object(launch.subprocess, 'Popen', side_effect=children), \ + patch.object(launch.uvicorn, 'Server') as server: + server.return_value.run.side_effect = RuntimeError('startup failed') + with self.assertRaisesRegex(RuntimeError, 'startup failed'): + launch.main() + children[0].terminate.assert_called_once() + + def test_stubborn_child_is_killed_and_reaped(self): + child = self.children()[0] + child.wait.side_effect = [subprocess.TimeoutExpired('child', 5), 0] + launch.stop_child(child) + child.kill.assert_called_once() + self.assertEqual(child.wait.call_count, 2) + + def test_exited_child_needs_no_signal(self): + child = Mock(poll=Mock(return_value=0)) + launch.stop_child(child) + child.terminate.assert_not_called() + + def test_sigterm_exits_through_child_cleanup(self): + children = self.children() + previous = signal.getsignal(signal.SIGTERM) + with patch.object(launch.subprocess, 'Popen', side_effect=children), \ + patch.object(launch.uvicorn, 'Server') as server: + server.return_value.run.side_effect = lambda **kwargs: signal.raise_signal(signal.SIGTERM) + with self.assertRaises(SystemExit) as caught: + launch.main() + self.assertEqual(caught.exception.code, 128 + signal.SIGTERM) + self.assertEqual(signal.getsignal(signal.SIGTERM), previous) + for child in children: + child.terminate.assert_called_once() + + +class PreboundUvicornTest(unittest.TestCase): + def test_http_is_ready_only_after_lifespan(self): + entered = threading.Event() + release = threading.Event() + + async def app(scope, receive, send): + if scope['type'] == 'lifespan': + await receive() + entered.set() + while not release.is_set(): + await asyncio.sleep(0.01) + await send({'type': 'lifespan.startup.complete'}) + await receive() + await send({'type': 'lifespan.shutdown.complete'}) + else: + await send({'type': 'http.response.start', 'status': 200, 'headers': []}) + await send({'type': 'http.response.body', 'body': b'ready'}) + + with socket.socket() as listener: + listener.bind(('127.0.0.1', 0)) + listener.listen(socket.SOMAXCONN) + server = launch.uvicorn.Server(launch.uvicorn.Config(app, lifespan='on', log_level='error')) + worker = threading.Thread(target=server.run, kwargs={'sockets': [listener]}, daemon=True) + worker.start() + try: + self.assertTrue(entered.wait(3), 'Uvicorn did not enter startup') + with socket.create_connection(listener.getsockname(), timeout=1) as client: + client.sendall(b'GET /health HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n') + client.settimeout(0.1) + with self.assertRaises(socket.timeout): + client.recv(4096) + release.set() + client.settimeout(3) + response = b'' + while True: + data = client.recv(4096) + if not data: + break + response += data + self.assertIn(b'200 OK', response) + self.assertIn(b'ready', response) + finally: + release.set() + server.should_exit = True + worker.join(timeout=5) + self.assertFalse(worker.is_alive(), 'Uvicorn shutdown did not complete') diff --git a/examples/run-code-env-inheritance/tests/test_startup.py b/examples/run-code-env-inheritance/tests/test_startup.py new file mode 100644 index 0000000000..4b325828c3 --- /dev/null +++ b/examples/run-code-env-inheritance/tests/test_startup.py @@ -0,0 +1,129 @@ +import unittest +import importlib +import os +from unittest.mock import AsyncMock, patch + +import httpx +import server + + +class ContextStartupTest(unittest.IsolatedAsyncioTestCase): + async def test_empty_envd_port_uses_default_readiness_endpoint(self): + with patch.dict(os.environ, {'ENVD_PORT': ''}): + importlib.reload(server) + self.assertEqual(server.ENVD_ENVS_URL, 'http://127.0.0.1:49983/envs') + importlib.reload(server) + + async def test_envd_unsupported_endpoint_fails_without_readiness_retries(self): + calls = [] + def respond(request): + calls.append(request) + return httpx.Response(404) + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client: + with self.assertRaisesRegex(RuntimeError, '404'): + await server._wait_for_envd(client) + self.assertEqual(len(calls), 1) + + async def test_envd_invalid_payload_fails_without_readiness_retries(self): + async with httpx.AsyncClient(transport=httpx.MockTransport(lambda request: httpx.Response(200, json=[]))) as client: + with self.assertRaisesRegex(RuntimeError, 'invalid environment payload'): + await server._wait_for_envd(client) + + async def test_envd_transient_failure_can_become_ready(self): + responses = iter([httpx.Response(503), httpx.Response(200, json={})]) + async with httpx.AsyncClient(transport=httpx.MockTransport(lambda request: next(responses))) as client: + await server._wait_for_envd(client) + + async def test_failed_connection_cleanup_still_deletes_known_session(self): + client = AsyncMock() + client.post.return_value = httpx.Response(200, json={"kernel": {"id": "kernel"}, "id": "session"}) + kernel = AsyncMock() + kernel.connect.side_effect = TimeoutError("handshake timed out") + kernel.close.side_effect = OSError("socket close failed") + with patch.object(server, "client", client), \ + patch.object(server, "_resolve_cwd", return_value="/workspace"), \ + patch.object(server, "KernelContext", return_value=kernel), \ + patch.object(server, "contexts", {}): + with self.assertRaisesRegex(TimeoutError, "handshake timed out"): + await server.create_context() + client.delete.assert_awaited_once_with(f"{server.JUPYTER_BASE_URL}/api/sessions/session") + + async def test_kernel_handshake_has_explicit_timeout(self): + kernel = server.KernelContext("kernel", "session", "/workspace") + ws = AsyncMock() + with patch.object(server, "connect", AsyncMock(return_value=ws)) as connect, \ + patch.object(server, "CONTEXT_CREATE_TIMEOUT", 75), \ + patch.object(kernel, "_receive_messages", AsyncMock()), \ + patch.object(kernel, "_run_background", AsyncMock()): + await kernel.connect() + self.assertEqual(connect.await_args.kwargs["open_timeout"], 75) + await kernel.close() + + async def test_failed_handshake_cleans_known_session(self): + client = AsyncMock() + client.post.return_value = httpx.Response(200, json={"kernel": {"id": "kernel"}, "id": "session"}) + kernel = AsyncMock() + kernel.connect.side_effect = TimeoutError("handshake timed out") + with patch.object(server, "client", client), \ + patch.object(server, "_resolve_cwd", return_value="/workspace"), \ + patch.object(server, "KernelContext", return_value=kernel), \ + patch.object(server, "contexts", {}) as contexts: + with self.assertRaises(TimeoutError): + await server.create_context() + self.assertEqual(contexts, {}) + kernel.close.assert_awaited_once() + client.delete.assert_awaited_once_with(f"{server.JUPYTER_BASE_URL}/api/sessions/session") + client.post.assert_awaited_once() + + async def test_context_creation_uses_explicit_timeout(self): + client = AsyncMock() + client.post.return_value = httpx.Response( + 200, json={"kernel": {"id": "kernel"}, "id": "session"} + ) + kernel = AsyncMock() + with patch.object(server, "client", client), \ + patch.object(server, "CONTEXT_CREATE_TIMEOUT", 75), \ + patch.object(server, "_resolve_cwd", return_value="/workspace"), \ + patch.object(server, "KernelContext", return_value=kernel), \ + patch.object(server, "contexts", {}): + result = await server.create_context() + self.assertEqual(result.id, "kernel") + self.assertEqual(client.post.await_args.kwargs["timeout"], 75) + client.post.assert_awaited_once() + kernel.connect.assert_awaited_once() + + async def test_timeout_does_not_repeat_session_creation(self): + client = AsyncMock() + client.post.side_effect = httpx.ReadTimeout("kernel startup slow") + with patch.object(server, "client", client), \ + patch.object(server, "_resolve_cwd", return_value="/workspace"): + with self.assertRaises(httpx.ReadTimeout): + await server.create_context() + client.post.assert_awaited_once() + + async def test_startup_failure_closes_http_client(self): + client = AsyncMock() + with patch.object(server, "_wait_for_jupyter", AsyncMock()), \ + patch.object(server, "_wait_for_envd", AsyncMock()), \ + patch.object(server.httpx, "AsyncClient", return_value=client), \ + patch.object(server, "contexts", {}), \ + patch.object(server, "create_context", AsyncMock(side_effect=httpx.ReadTimeout("slow"))): + with self.assertRaises(httpx.ReadTimeout): + async with server.lifespan(server.app): + self.fail("failed startup must not become ready") + client.aclose.assert_awaited_once() + self.assertIsNone(server.default_context_id) + + async def test_missing_envd_never_reports_interpreter_ready(self): + client = AsyncMock() + create = AsyncMock() + with patch.object(server, "_wait_for_jupyter", AsyncMock()), \ + patch.object(server, "_wait_for_envd", AsyncMock(side_effect=RuntimeError("envd did not become ready"))), \ + patch.object(server.httpx, "AsyncClient", return_value=client), \ + patch.object(server, "contexts", {}), \ + patch.object(server, "create_context", create): + with self.assertRaisesRegex(RuntimeError, "envd did not become ready"): + async with server.lifespan(server.app): + self.fail("missing envd must not become ready") + create.assert_not_awaited() + client.aclose.assert_awaited_once()