diff --git a/docs/remote_job_monitoring.md b/docs/remote_job_monitoring.md index f96fce1d..b9a6df5b 100644 --- a/docs/remote_job_monitoring.md +++ b/docs/remote_job_monitoring.md @@ -1,19 +1,32 @@ # Remote Job Monitoring -MatCreator manages E2B sandboxes as durable, session-scoped remote jobs. The -remote-job control plane separates a sandbox's provider identity and liveness -from the agent step that created it, so the FastAPI frontend can observe and -control the sandbox after an agent, browser, or middleware request reconnects. +MatCreator manages sandboxes and batch jobs as durable, session-scoped remote +jobs. The remote-job control plane separates a job's provider identity and +liveness from the agent step that created it, so the FastAPI frontend can +observe and control it after an agent, browser, or middleware request +reconnects. + +Every provider-specific operation goes through a small adapter protocol (see +[Provider Plugin Architecture](#provider-plugin-architecture) below), so +`RemoteJobService`, `RemoteJobMonitor`, and the web API never branch on a +provider name. Built in providers today: `e2b` (interactive sandbox via the +E2B SDK), `bohr_sandbox` (interactive sandbox via the `bohr` CLI), and +`bohr_job` (batch/HPC-style job via `bohr job submit`). ## Architecture ```mermaid flowchart LR - Agent[Step executor] --> Tools[E2B tools] + Agent[Step executor] --> Tools[remote_job_tools] Tools --> Service[RemoteJobService] Service --> Store[(remote-jobs.db)] - Service --> Adapter[E2BSandboxAdapter] - Adapter --> Sandbox[E2B/Bohrium sandbox] + Service --> Registry[providers registry] + Registry --> E2B[E2BSandboxAdapter] + Registry --> BohrSbx[BohrSandboxAdapter] + Registry --> BohrJob[BohrJobAdapter] + E2B --> Sandbox[E2B/Bohrium sandbox] + BohrSbx --> Sandbox + BohrJob --> Batch[Bohrium batch job] Monitor[RemoteJobMonitor] --> Service Monitor --> Store @@ -23,7 +36,7 @@ flowchart LR ``` The SQLite record is the source of truth for MatCreator's normalized job -lifecycle. The provider sandbox remains the source of truth for provider +lifecycle. The provider job/sandbox remains the source of truth for provider liveness. This distinction lets the UI report both a meaningful lifecycle state and the latest connectivity observation without conflating them. @@ -32,28 +45,38 @@ and the latest connectivity observation without conflating them. | Component | Location | Responsibility | | --- | --- | --- | | `RemoteJobStore` | `src/matcreator/control_plane/remote_jobs.py` | Persists jobs, lifecycle transitions, provider snapshots, and user-control events in SQLite. | -| `RemoteJobService` | `src/matcreator/control_plane/remote_job_service.py` | Coordinates E2B operations with durable records and enforces valid lifecycle operations. | -| `E2BSandboxAdapter` | `src/matcreator/control_plane/e2b.py` | Small lazy-import boundary around the E2B SDK for sandbox creation, commands, files, pause, kill, and probing. | -| `RemoteJobMonitor` | `src/matcreator/control_plane/remote_job_monitor.py` | Periodically reconciles active E2B records and applies bounded backoff after failed probes. | -| Agent tools | `src/matcreator/agents/execution_agent/e2b_tools.py` | Submit and operate on jobs owned by the current session. | -| Middleware APIs | `web/main.py` | List jobs/events and offer session-owner pause, terminate, and refresh endpoints. | +| `RemoteJobService` | `src/matcreator/control_plane/remote_job_service.py` | Coordinates provider operations with durable records and enforces valid lifecycle operations, dispatching to the adapter registered for each job's `provider`. | +| `RemoteJobAdapter` protocol | `src/matcreator/control_plane/providers/base.py` | The boundary every provider implements: `create`/`status`/`cancel` are mandatory; `pause`/`resume`/`run_command`/`upload_file`/`download_file`/`collect_outputs` are gated by declared `RemoteJobCapability` flags. | +| Provider registry | `src/matcreator/control_plane/providers/registry.py` | Maps a provider name to a lazily constructed adapter instance. | +| `E2BSandboxAdapter` | `src/matcreator/control_plane/providers/e2b.py` | Interactive sandbox via the E2B SDK: create, commands, files, pause, kill, probe. | +| `BohrSandboxAdapter` | `src/matcreator/control_plane/providers/bohr_sandbox.py` | Interactive sandbox via the `bohr` CLI (`bohr sandbox create/exec/files/describe/delete`). No pause/resume — the CLI has no such subcommand. | +| `BohrJobAdapter` | `src/matcreator/control_plane/providers/bohr_job.py` | Batch/HPC-style job via the `bohr` CLI (`bohr job submit/describe/download/terminate`). Submit-time inputs only; no interactive exec. | +| `RemoteJobMonitor` | `src/matcreator/control_plane/remote_job_monitor.py` | Periodically reconciles active jobs of every registered provider, using each adapter's own `poll_interval_seconds` for backoff scheduling. | +| Agent tools | `src/matcreator/agents/execution_agent/remote_job_tools.py` | Provider-specific submit tools (`submit_e2b_sandbox`, `submit_bohr_sandbox`, `submit_bohr_job`) plus provider-generic post-submission tools that dispatch on `job_id` alone. | +| Middleware APIs | `web/main.py` | List jobs/events and offer session-owner pause, terminate, and refresh endpoints, generic across providers. | ## Submission and Persistence -`submit_e2b_sandbox` requires an explicit template. It creates a deterministic -idempotency key from the session, execution node, and template, then delegates -to `RemoteJobService.submit_e2b`. +Submission is provider-specific — an interactive sandbox needs a template +while a batch job needs a machine type and image — so there is one submit +tool per provider: `submit_e2b_sandbox`, `submit_bohr_sandbox`, +`submit_bohr_job`. Each builds a deterministic idempotency key from the +session, execution node, and a provider-specific discriminator, then +delegates to `RemoteJobService.submit_job(provider=..., spec=...)`. The service creates the SQLite job record before making the provider request. -The persisted specification contains the template, endpoint, project ID, -timeout, lifecycle policy, and metadata, but never the API key. Repeated calls -with the same idempotency key return the existing job instead of creating a -second sandbox. - -Once sandbox creation succeeds, the service stores the provider sandbox ID in -`external_id` and transitions the job to `running`. Agent recovery records the -job reference against the execution graph so an interrupted execution can wait -for or accurately report an existing sandbox rather than resubmitting it. +`persisted_specification` — everything in `spec` except secrets like an API +key — is what actually gets stored; `spec` itself (which may contain +secrets) is passed to the adapter's `create` but never persisted. Repeated +calls with the same idempotency key return the existing job instead of +creating a second sandbox or job. + +Once creation succeeds, the service stores the provider-side ID in +`external_id`, probes the adapter once for an initial status (letting a batch +provider start in `queued` instead of always assuming `running`), and +transitions the job accordingly. Agent recovery records the job reference +against the execution graph so an interrupted execution can wait for or +accurately report an existing job rather than resubmitting it. ## Lifecycle and Observations @@ -62,7 +85,8 @@ Important normalized states include: - `created`, `submitting`, `queued`, `running`, `paused`, and `resuming` for active work. -- `succeeded` and `collecting` while a job's results are being handled. +- `succeeded` and `collecting` while a batch job's results are being pulled + via `collect_remote_job_outputs`. - `collected`, `failed`, `cancelled`, `terminated`, and `lost` as terminal outcomes. @@ -71,16 +95,27 @@ transitions use optimistic concurrency checks, so stale pause, terminate, or provider updates cannot silently overwrite newer state. Provider probe data is stored in `snapshot`; examples include -`provider_status`, `sandbox_id`, `last_command_exit_code`, and `last_upload`. -An observation does not itself alter the normalized lifecycle state. +`provider_status`, `sandbox_id`, `phase` (for a batch job), `last_command_exit_code`, +and `last_upload`. An observation does not itself alter the normalized +lifecycle state unless the adapter reports a `normalized_status` that differs +from the current one — see [Provider Plugin Architecture](#provider-plugin-architecture). ## Monitoring and Refresh -`RemoteJobMonitor` considers active E2B jobs and probes jobs in `queued`, -`running`, `submitting`, or `resuming` states. A successful probe records a -reachable provider snapshot. A failed probe records `provider_status` as -`unreachable` and increases the next probe delay exponentially, bounded by the -configured maximum backoff. +`RemoteJobMonitor` considers active jobs of every registered provider and +probes jobs in `queued`, `running`, `submitting`, or `resuming` states, using +each job's own adapter to decide how — and how often — to probe. A batch +provider like `bohr_job` declares a much longer `poll_interval_seconds` (60s) +than an interactive sandbox (15s), so it is polled far less often without any +special-casing in the monitor itself. + +For an interactive adapter (`e2b`, `bohr_sandbox`) a successful probe records +a reachable provider snapshot; a failed probe records `provider_status` as +`unreachable` and increases the next probe delay exponentially, bounded by +the configured maximum backoff. For a batch adapter (`bohr_job`) the same +probe can report a `normalized_status` change (e.g. `queued` -> `running` -> +`succeeded`/`failed`/`cancelled`), which the service turns into an actual +lifecycle transition instead of just an observation. Monitor schedules are intentionally in memory. The job records themselves are durable, so a restarted monitor begins by reconciling active jobs from SQLite. @@ -129,11 +164,11 @@ mechanism rather than two. Re-attachment is explicit rather than accidental. When a node that already owns a job runs again, the runner injects the job's identity into the executor's -`prior_context` with instructions to call `get_e2b_job_status` and never call -`submit_e2b_sandbox` for that step. In Flash mode, which has no execution graph, -a step's node ID is derived from its label or a hash of its action, so a repeated -step keeps the same submission idempotency key and re-attaches instead of -creating a duplicate sandbox. +`prior_context` with instructions to call `get_remote_job_status` and never +call any of the `submit_*` tools for that step. In Flash mode, which has no +execution graph, a step's node ID is derived from its label or a hash of its +action, so a repeated step keeps the same submission idempotency key and +re-attaches instead of creating a duplicate job. ## Controls and Ownership @@ -147,11 +182,13 @@ POST /api/sessions/{session_id}/remote-jobs/{job_id}/terminate Both invoke the provider operation through `RemoteJobService`, update the durable job lifecycle, and append a `user_control` event. They do not cancel the step-executor process. The executor sees this event through -`get_e2b_job_status` and must report `needs_replanning` rather than retrying an -interrupted command or submitting a replacement sandbox. +`get_remote_job_status` and must report `needs_replanning` rather than +retrying an interrupted command or submitting a replacement job. `pause` +returns a 409 (via `CapabilityError`) for a provider that does not support +pausing, such as `bohr_job`. -`terminate_e2b_sandbox` irreversibly releases a sandbox. Agents should collect -or record required output before calling it. +`terminate_remote_job` irreversibly releases a job or sandbox. Agents should +collect or record required output before calling it. ## Storage Scope @@ -160,12 +197,72 @@ the middleware routes each owner to a per-user `.adk/remote-jobs.db` under the user's mounted MatCreator home. This keeps job records, controls, and monitoring isolated by owner and session. +## Provider Plugin Architecture + +Adding a new remote-job provider (a different HPC scheduler, another +sandbox platform, ...) means implementing `RemoteJobAdapter` and registering +it — nothing else in the control plane changes. + +1. **Implement the adapter** (`src/matcreator/control_plane/providers/.py`): + subclass `RemoteJobAdapter` from `providers/base.py` and implement the + three mandatory methods (`create`, `status`, `cancel`). Declare + `provider`, `capabilities` (a `frozenset[RemoteJobCapability]`), and + `poll_interval_seconds` as class attributes. Implement only the optional + methods your capabilities declare: + + | Capability | Optional method(s) | Example provider | + | --- | --- | --- | + | `PAUSE` / `RESUME` | `pause` / `resume` | `e2b` (pause only) | + | `INTERACTIVE_EXEC` | `run_command` | `e2b`, `bohr_sandbox` | + | `FILE_TRANSFER` | `upload_file` / `download_file` | `e2b`, `bohr_sandbox` | + | `BATCH_COLLECT` | `collect_outputs` | `bohr_job` | + + `status` returns a `RemoteJobStatus(normalized_status, snapshot, error)`. + Use `normalized_status=None` when the provider can only confirm liveness + (an interactive sandbox that stays "running" until explicitly stopped); + return one of the canonical statuses from `remote_jobs.py` (e.g. + `"succeeded"`, `"failed"`, `"cancelled"`) when the provider can report an + actual lifecycle observation (a batch job that finishes on its own). + +2. **Register it** in `src/matcreator/control_plane/providers/__init__.py` + with a lazy factory: + ```python + register_adapter("my_provider", lambda: MyProviderAdapter()) + ``` + The factory is not called until the first `get_adapter("my_provider")`, so + registering a provider never forces an optional SDK/CLI import at process + startup. + +3. **(Optional) add a submit tool** in + `src/matcreator/agents/execution_agent/remote_job_tools.py` if the agent + should be able to submit this provider's jobs — submission parameters are + inherently provider-specific (a template vs. a machine type + image), so + this is the one place a new provider needs new code beyond the adapter + itself. Every operation *after* submission + (`get_remote_job_status`/`pause_remote_job`/`terminate_remote_job`/ + `run_remote_job_command`/`upload_remote_job_input`/ + `download_remote_job_output`/`collect_remote_job_outputs`) already works + for any provider without changes, dispatching on the stored `job_id` alone. + +`RemoteJobService` and `RemoteJobMonitor` never import a specific adapter or +branch on a provider name — they resolve the adapter for a job through the +registry (`RemoteJobService.adapter_for`) and check `adapter.capabilities` +before calling an optional method, raising `CapabilityError` with a clear, +provider-attributed message if unsupported (e.g. pausing a `bohr_job`). + ## Operational Notes -- The control plane currently supports E2B sandboxes, although the persistent - store is provider-neutral by design. +- Built-in providers: `e2b` (interactive, via the E2B SDK), `bohr_sandbox` + (interactive, via the `bohr` CLI), and `bohr_job` (batch/HPC-style, via the + `bohr` CLI). The persistent store and service are provider-neutral by + design; see [Provider Plugin Architecture](#provider-plugin-architecture) + to add another. - Commands do not persist command text or output in the remote-job database; only limited operational telemetry is recorded. - A sandbox's configured creation timeout is distinct from the monitoring - interval. The adapter currently passes `timeout=0` to E2B command execution, - leaving command duration unrestricted by this control plane. + interval. The E2B adapter currently passes `timeout=0` to command + execution, leaving command duration unrestricted by this control plane. +- `bohr_job` only supports single-job submission (`bohr job submit`); `bohr + job_group` fan-out (many jobs sharing one group) is a possible future + adapter, not implemented here. + diff --git a/pyproject.toml b/pyproject.toml index a70e5f82..eb6b57be 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,7 @@ dependencies = [ "phonopy>=2.38.0", "seekpath>=2.1.0", "rdkit>=2025.9.1", - "google-adk>=1.28.0", + "google-adk>=2.0.0", "google-cloud-storage", "lbg>=1.2.29", "litellm>=1.77.4", @@ -29,6 +29,7 @@ dependencies = [ "mp-api", "bcrypt>=4.0.0", "docker>=7.0.0", + "e2b<=2.20.0", "e2b-code-interpreter", "matcraft-kit", "know-do-graph>=0.1.8", diff --git a/src/matcreator/agents/execution_agent/e2b_tools.py b/src/matcreator/agents/execution_agent/e2b_tools.py deleted file mode 100644 index 79fea5de..00000000 --- a/src/matcreator/agents/execution_agent/e2b_tools.py +++ /dev/null @@ -1,232 +0,0 @@ -"""Tracked E2B sandbox tools available to isolated step executors.""" -from __future__ import annotations - -import hashlib -import os -from pathlib import Path -from typing import Any - -from google.adk.tools.tool_context import ToolContext - -from ...control_plane.remote_job_service import E2BConnectionConfig, RemoteJobService -from ...control_plane.remote_jobs import RemoteJobStore -from ...workspace import ADK_DIR -from .recovery import record_remote_job_reference - - -def _service() -> RemoteJobService: - return RemoteJobService(RemoteJobStore(ADK_DIR / "remote-jobs.db")) - - -def _owner_id(tool_context: ToolContext) -> str: - invocation = getattr(tool_context, "_invocation_context", None) - return str(getattr(invocation, "user_id", "") or tool_context.state.get("user_id") or "default") - - -def _node_id(tool_context: ToolContext) -> str: - graph_node = str(tool_context.state.get("_graph_exec_node_id") or "step") - return graph_node.rsplit("__node_", 1)[-1] - - -def _connection() -> E2BConnectionConfig: - # Bohrium E2B endpoint uses bare hex keys; disable SDK format validation - os.environ.setdefault("E2B_VALIDATE_API_KEY", "false") - return E2BConnectionConfig( - api_key=os.environ.get("E2B_API_KEY", ""), - api_url=os.environ.get("E2B_API_URL", ""), - project_id=os.environ.get("BOHRIUM_PROJECT_ID", ""), - template="", - ) - - -def submit_e2b_sandbox( - tool_context: ToolContext, - *, - timeout: int = 7200, - template: str = None, - lifecycle: dict[str, Any] | None = None, -) -> dict[str, Any]: - """Create or reuse a tracked E2B sandbox for the current execution step. - - The configured E2B API key, endpoint, and project ID are used server-side. - Never use shell commands or include credentials in tool inputs. A repeated - call for the same step and template returns the existing sandbox record. - """ - session_id = str(tool_context.state.get("session_id") or "") - if not session_id: - return {"status": "error", "message": "No session_id is available for E2B submission."} - node_id = _node_id(tool_context) - connection = _connection() - if not template: - return { - "status": "error", - "message": "An explicit E2B sandbox template is required. Use 'lbg sdbx template ls -q' to list available templates.", - } - connection = E2BConnectionConfig( - api_key=connection.api_key, - api_url=connection.api_url, - project_id=connection.project_id, - template=template, - ) - identity = f"{session_id}:{node_id}:{connection.template}" - idempotency_key = f"e2b:{hashlib.sha256(identity.encode()).hexdigest()}" - try: - job = _service().submit_e2b( - owner_id=_owner_id(tool_context), - session_id=session_id, - node_id=node_id, - step_number=tool_context.state.get("step_number"), - idempotency_key=idempotency_key, - connection=connection, - timeout=timeout, - lifecycle=lifecycle or {"on_timeout": "pause", "auto_resume": True}, - ) - except Exception as exc: - return {"status": "error", "message": f"E2B submission failed: {exc}"} - record_remote_job_reference( - session_id=session_id, - node_id=node_id, - job_id=job["job_id"], - provider="e2b", - external_id=job["external_id"], - ) - return { - "status": job["status"], - "job_id": job["job_id"], - "sandbox_id": job["external_id"], - "message": "Tracked E2B sandbox is ready. Use its job_id for status or controls.", - } - - -def get_e2b_job_status(job_id: str, tool_context: ToolContext) -> dict[str, Any]: - """Read one tracked E2B job owned by the current session.""" - service = _service() - job = service.store.get_job(job_id) - if job is None or job["owner_id"] != _owner_id(tool_context) or job["session_id"] != tool_context.state.get("session_id"): - return {"status": "error", "message": "E2B job was not found in this session."} - result = {key: job[key] for key in ("job_id", "status", "external_id", "snapshot", "error", "updated_at")} - controls = [ - event["payload"] - for event in service.store.list_events(job_id) - if event["event_type"] == "user_control" - ] - if controls: - result["user_control"] = controls[-1] - return result - - -def pause_e2b_sandbox(job_id: str, tool_context: ToolContext) -> dict[str, Any]: - """Pause a tracked E2B sandbox belonging to the current session.""" - job = get_e2b_job_status(job_id, tool_context) - if job.get("status") == "error": - return job - try: - paused = _service().pause_e2b(job_id) - except Exception as exc: - return {"status": "error", "message": f"E2B pause failed: {exc}"} - return {"job_id": paused["job_id"], "status": paused["status"], "sandbox_id": paused["external_id"]} - - -def terminate_e2b_sandbox(job_id: str, tool_context: ToolContext) -> dict[str, Any]: - """Terminate a tracked E2B sandbox belonging to the current session.""" - job = get_e2b_job_status(job_id, tool_context) - if job.get("status") == "error": - return job - try: - terminated = _service().terminate_e2b(job_id) - except Exception as exc: - return {"status": "error", "message": f"E2B termination failed: {exc}"} - return {"job_id": terminated["job_id"], "status": terminated["status"], "sandbox_id": terminated["external_id"]} - - -def run_e2b_command( - job_id: str, - command: str, - tool_context: ToolContext, - user: str = "root", -) -> dict[str, Any]: - """Run one command inside a tracked E2B sandbox in the current session. - - Do not put credentials in ``command``. Command text and output are returned - to the current step but are not persisted in the durable job snapshot. - """ - job = get_e2b_job_status(job_id, tool_context) - if job.get("status") == "error": - return job - try: - return _service().run_e2b_command(job_id, command, user=user) - except Exception as exc: - current = get_e2b_job_status(job_id, tool_context) - result = {"status": "error", "message": f"E2B command failed: {exc}"} - if current.get("user_control"): - result["user_control"] = current["user_control"] - return result - - -def _resolve_workspace_child( - tool_context: ToolContext, - user_path: str, -) -> tuple[Path | None, str | None]: - """Resolve ``user_path`` against the current workspace, confining it. - - Returns ``(resolved_path, None)`` on success or ``(None, message)`` if the - workspace is unavailable or the path escapes it. Shared by upload (source) - and download (destination) so confinement logic cannot drift between them. - """ - workspace_dir = tool_context.state.get("workspace_dir") - if not workspace_dir: - return None, "No workspace_dir is available for the current step." - workspace = Path(str(workspace_dir)).resolve() - candidate = Path(user_path).expanduser() - candidate = candidate.resolve() if candidate.is_absolute() else (workspace / candidate).resolve() - if not candidate.is_relative_to(workspace): - return None, "Path must resolve inside the current workspace." - return candidate, None - - -def upload_e2b_input( - job_id: str, - source_path: str, - destination_path: str, - tool_context: ToolContext, -) -> dict[str, Any]: - """Upload a workspace input file into a tracked E2B sandbox. - - ``source_path`` must resolve inside the current workspace. Use an absolute - sandbox path for ``destination_path`` such as ``/home/user/input.in``. - """ - job = get_e2b_job_status(job_id, tool_context) - if job.get("status") == "error": - return job - source, error = _resolve_workspace_child(tool_context, source_path) - if error is not None: - return {"status": "error", "message": f"E2B upload failed: {error}"} - try: - return _service().upload_e2b_file(job_id, source, destination_path) - except Exception as exc: - return {"status": "error", "message": f"E2B upload failed: {exc}"} - - -def download_e2b_output( - job_id: str, - source_path: str, - destination_path: str, - tool_context: ToolContext, -) -> dict[str, Any]: - """Download a file from a tracked E2B sandbox into the local workspace. - - ``source_path`` is an absolute path inside the sandbox (e.g. - ``/home/user/CHGCAR``). ``destination_path`` must resolve inside the - current workspace. Large binary outputs are streamed via the E2B - filesystem API, so they are not truncated by command-output limits. - """ - job = get_e2b_job_status(job_id, tool_context) - if job.get("status") == "error": - return job - destination, error = _resolve_workspace_child(tool_context, destination_path) - if error is not None: - return {"status": "error", "message": f"E2B download failed: {error}"} - try: - return _service().download_e2b_file(job_id, source_path, destination) - except Exception as exc: - return {"status": "error", "message": f"E2B download failed: {exc}"} \ No newline at end of file diff --git a/src/matcreator/agents/execution_agent/recovery.py b/src/matcreator/agents/execution_agent/recovery.py index 5329b83a..1cab4242 100644 --- a/src/matcreator/agents/execution_agent/recovery.py +++ b/src/matcreator/agents/execution_agent/recovery.py @@ -286,11 +286,16 @@ def _active_remote_job( def remote_job_reference(job: dict[str, Any]) -> dict[str, Any]: """Return the identity subset of a remote job stored on a graph node.""" + snapshot = job.get("snapshot") if isinstance(job.get("snapshot"), dict) else {} return { "job_id": job["job_id"], "provider": job["provider"], "external_id": job["external_id"], "status": job["status"], + # Surfaced so re-attach instructions can tell a fresh executor to poll + # an in-flight background command instead of guessing whether one is + # running or safely re-issuing it (see start_job_command/poll_job_command). + "has_background_command": bool(snapshot.get("background_command")), } diff --git a/src/matcreator/agents/execution_agent/remote_job_tools.py b/src/matcreator/agents/execution_agent/remote_job_tools.py new file mode 100644 index 00000000..06d80a5f --- /dev/null +++ b/src/matcreator/agents/execution_agent/remote_job_tools.py @@ -0,0 +1,545 @@ +"""Remote-job tools available to isolated step executors. + +Submission is provider-specific — an E2B/bohr sandbox needs a template while +a batch job needs a machine type and image, so there is one submit tool per +provider (``submit_e2b_sandbox``, ``submit_bohr_sandbox``, +``submit_bohr_job``). Every operation after submission dispatches on the +``job_id`` alone and works the same for any provider, so adding a new +provider plugin never requires a new post-submission tool here. +""" +from __future__ import annotations + +import hashlib +import json +import os +from pathlib import Path +from typing import Any + +from google.adk.tools.tool_context import ToolContext + +from ...control_plane.providers.e2b import E2BConnectionConfig +from ...control_plane.remote_job_service import RemoteJobService +from ...control_plane.remote_jobs import TERMINAL_REMOTE_JOB_STATUSES, RemoteJobStore +from ...workspace import ADK_DIR +from .recovery import record_remote_job_reference + +# Every terminal status except "collected" (the successful end of a batch +# job) means the submission is not usable and must not be reported as ready. +_FAILED_SUBMISSION_STATUSES = TERMINAL_REMOTE_JOB_STATUSES - {"collected"} + + +def _service() -> RemoteJobService: + return RemoteJobService(RemoteJobStore(ADK_DIR / "remote-jobs.db")) + + +def _owner_id(tool_context: ToolContext) -> str: + invocation = getattr(tool_context, "_invocation_context", None) + return str(getattr(invocation, "user_id", "") or tool_context.state.get("user_id") or "default") + + +def _node_id(tool_context: ToolContext) -> str: + graph_node = str(tool_context.state.get("_graph_exec_node_id") or "step") + return graph_node.rsplit("__node_", 1)[-1] + + +def _idempotency_key(session_id: str, node_id: str, discriminator: str) -> str: + identity = f"{session_id}:{node_id}:{discriminator}" + return f"remote-job:{hashlib.sha256(identity.encode()).hexdigest()}" + + +def _submit( + tool_context: ToolContext, + *, + provider: str, + spec: dict[str, Any], + discriminator: str, + persisted_specification: dict[str, Any] | None = None, +) -> dict[str, Any]: + """Shared submission plumbing used by every provider-specific submit tool.""" + session_id = str(tool_context.state.get("session_id") or "") + if not session_id: + return {"status": "error", "message": "No session_id is available for remote-job submission."} + node_id = _node_id(tool_context) + idempotency_key = _idempotency_key(session_id, node_id, discriminator) + try: + job = _service().submit_job( + owner_id=_owner_id(tool_context), + session_id=session_id, + provider=provider, + node_id=node_id, + step_number=tool_context.state.get("step_number"), + idempotency_key=idempotency_key, + spec=spec, + persisted_specification=persisted_specification, + ) + except Exception as exc: + return {"status": "error", "message": f"{provider} submission failed: {exc}"} + record_remote_job_reference( + session_id=session_id, + node_id=node_id, + job_id=job["job_id"], + provider=provider, + external_id=job["external_id"], + ) + return { + "status": job["status"], + "job_id": job["job_id"], + "external_id": job["external_id"], + "error": job.get("error"), + } + + +def _submission_response(result: dict[str, Any], *, id_field: str, success_message: str) -> dict[str, Any]: + """Convert a ``_submit`` result into the tool response, never claiming a + + failed/cancelled/terminated/lost job is ready. The durable record's + ``error`` is surfaced so the caller sees the actual cause. + """ + if result.get("status") == "error": + return result + response = { + "status": result["status"], + "job_id": result["job_id"], + id_field: result["external_id"], + } + if result["status"] in _FAILED_SUBMISSION_STATUSES: + cause = result.get("error") or f"the tracked job is in terminal status '{result['status']}'" + response["message"] = f"Remote job submission is not usable: {cause}" + return response + response["message"] = success_message + return response + + +def _connection() -> E2BConnectionConfig: + # Bohrium E2B endpoint uses bare hex keys; disable SDK format validation + os.environ.setdefault("E2B_VALIDATE_API_KEY", "false") + return E2BConnectionConfig( + api_key=os.environ.get("E2B_API_KEY", ""), + api_url=os.environ.get("E2B_API_URL", ""), + project_id=os.environ.get("BOHRIUM_PROJECT_ID", ""), + template="", + ) + + +def submit_e2b_sandbox( + tool_context: ToolContext, + *, + timeout: int = 7200, + template: str = None, + lifecycle: dict[str, Any] | str | None = None, +) -> dict[str, Any]: + """Create or reuse a tracked E2B sandbox for the current execution step. + + The configured E2B API key, endpoint, and project ID are used server-side. + Never use shell commands or include credentials in tool inputs. A repeated + call for the same step and template returns the existing sandbox record. + """ + session_id = str(tool_context.state.get("session_id") or "") + if not session_id: + return {"status": "error", "message": "No session_id is available for E2B submission."} + if not template: + return { + "status": "error", + "message": "An explicit E2B sandbox template is required. Use 'lbg sdbx template ls -q' to list available templates.", + } + if isinstance(lifecycle, str): + try: + lifecycle = json.loads(lifecycle) + except json.JSONDecodeError: + pass # still a str; rejected below + if lifecycle is not None and not isinstance(lifecycle, dict): + return { + "status": "error", + "message": "lifecycle must be a JSON object such as {\"on_timeout\": \"pause\", \"auto_resume\": true}.", + } + connection = _connection() + missing_config = [ + name + for name, value in ( + ("E2B_API_KEY", connection.api_key), + ("E2B_API_URL", connection.api_url), + ("BOHRIUM_PROJECT_ID", connection.project_id), + ) + if not value + ] + if missing_config: + return { + "status": "error", + "message": ( + f"E2B is not configured on the server: {', '.join(missing_config)} unset. " + "If only the `bohr` CLI is available, use submit_bohr_sandbox instead." + ), + } + connection = E2BConnectionConfig( + api_key=connection.api_key, + api_url=connection.api_url, + project_id=connection.project_id, + template=template, + ) + spec = connection.to_spec_dict(timeout=timeout, lifecycle=lifecycle or {"on_timeout": "pause", "auto_resume": True}) + persisted_specification = {key: value for key, value in spec.items() if key != "api_key"} + result = _submit( + tool_context, + provider="e2b", + spec=spec, + discriminator=template, + persisted_specification=persisted_specification, + ) + return _submission_response( + result, + id_field="sandbox_id", + success_message="Tracked E2B sandbox is ready. Use its job_id for status or controls.", + ) + + +def submit_bohr_sandbox( + tool_context: ToolContext, + *, + project_id: int = None, + template: str = None, + timeout: int = None, + image: str = None, + gpu: str = None, + never_timeout: bool = False, + env: dict[str, str] | None = None, +) -> dict[str, Any]: + """Create or reuse a tracked Bohrium CLI sandbox (`bohr sandbox`) for the current step. + + Both this and `submit_e2b_sandbox` reach the same Bohrium sandbox + platform; use this one when only the `bohr` CLI (not the E2B SDK/API key) + is available in the current environment. An explicit ``template`` is + required (e.g. ``doc-compiler``). ``gpu`` selects a GPU shortcut template + (``4090``/``5090``/``l20``). Falls back to the `BOHRIUM_PROJECT_ID` + environment variable if ``project_id`` is omitted. + """ + resolved_project_id = project_id or os.environ.get("BOHRIUM_PROJECT_ID", "") + if not resolved_project_id: + return {"status": "error", "message": "An explicit project_id is required for a bohr sandbox."} + if not template: + return { + "status": "error", + "message": ( + "An explicit sandbox template is required (e.g. 'doc-compiler'). " + "Use 'bohr sandbox template list' to see available templates." + ), + } + spec = { + "project_id": resolved_project_id, + "template": template, + "timeout": timeout, + "image": image, + "gpu": gpu, + "never_timeout": never_timeout, + "env": env or {}, + } + result = _submit( + tool_context, + provider="bohr_sandbox", + spec=spec, + discriminator=template, + ) + return _submission_response( + result, + id_field="sandbox_id", + success_message="Tracked bohr sandbox is ready. Use its job_id for status or controls.", + ) + + +def submit_bohr_job( + tool_context: ToolContext, + *, + project_id: int = None, + job_name: str = None, + machine_type: str = None, + image_address: str = None, + command: str = None, + input_directory: str | None = None, + result_path: str | None = None, + max_run_time: int | None = None, +) -> dict[str, Any]: + """Submit a batch/HPC-style Bohrium job (`bohr job submit`) for the current step. + + This is a fire-and-forget batch submission, not an interactive sandbox: + inputs are staged once via ``input_directory`` and there is no + `run_remote_job_command` for this provider — the whole computation must + be expressed in ``command``. Poll `get_remote_job_status` until it + reports ``succeeded``, then call `collect_remote_job_outputs`. + """ + resolved_project_id = project_id or os.environ.get("BOHRIUM_PROJECT_ID", "") + missing = [ + name + for name, value in ( + ("project_id", resolved_project_id), + ("job_name", job_name), + ("machine_type", machine_type), + ("image_address", image_address), + ("command", command), + ) + if not value + ] + if missing: + return { + "status": "error", + "message": f"Missing required field(s) for bohr job submission: {', '.join(missing)}", + } + spec = { + "project_id": resolved_project_id, + "job_name": job_name, + "machine_type": machine_type, + "image_address": image_address, + "command": command, + "input_directory": input_directory, + "result_path": result_path, + "max_run_time": max_run_time, + } + result = _submit( + tool_context, + provider="bohr_job", + spec=spec, + discriminator=f"{job_name}:{machine_type}:{image_address}", + ) + return _submission_response( + result, + id_field="bohr_job_id", + success_message=( + "Tracked bohr batch job is submitted. Poll get_remote_job_status until it " + "reports succeeded, then call collect_remote_job_outputs." + ), + ) + + +def get_remote_job_status(job_id: str, tool_context: ToolContext) -> dict[str, Any]: + """Read one tracked remote job (any provider) owned by the current session.""" + service = _service() + job = service.store.get_job(job_id) + if ( + job is None + or job["owner_id"] != _owner_id(tool_context) + or job["session_id"] != tool_context.state.get("session_id") + ): + return {"status": "error", "message": "Remote job was not found in this session."} + result = { + key: job[key] for key in ("job_id", "provider", "status", "external_id", "snapshot", "error", "updated_at") + } + controls = [ + event["payload"] for event in service.store.list_events(job_id) if event["event_type"] == "user_control" + ] + if controls: + result["user_control"] = controls[-1] + return result + + +def pause_remote_job(job_id: str, tool_context: ToolContext) -> dict[str, Any]: + """Pause a tracked remote job belonging to the current session. + + Returns an error if the job's provider does not support pausing (e.g. a + batch job); terminate it instead if it must stop. + """ + job = get_remote_job_status(job_id, tool_context) + if job.get("status") == "error": + return job + try: + paused = _service().pause_job(job_id) + except Exception as exc: + return {"status": "error", "message": f"Pause failed: {exc}"} + return {"job_id": paused["job_id"], "status": paused["status"], "external_id": paused["external_id"]} + + +def terminate_remote_job(job_id: str, tool_context: ToolContext) -> dict[str, Any]: + """Terminate a tracked remote job belonging to the current session.""" + job = get_remote_job_status(job_id, tool_context) + if job.get("status") == "error": + return job + try: + terminated = _service().terminate_job(job_id) + except Exception as exc: + return {"status": "error", "message": f"Termination failed: {exc}"} + return {"job_id": terminated["job_id"], "status": terminated["status"], "external_id": terminated["external_id"]} + + +def run_remote_job_command( + job_id: str, + command: str, + tool_context: ToolContext, + user: str = "root", +) -> dict[str, Any]: + """Run one short command inside a tracked interactive remote job (e.g. a sandbox). + + This BLOCKS until the command finishes, with no timeout of its own. Only + use it for commands expected to finish in well under a minute (checking a + file, `mkdir`, `grep`, listing a directory, ...). For anything that might + run longer — a training run, a `vasp_std`/`mpirun` invocation, any real + computation — use `start_remote_job_command` + `poll_remote_job_command` + instead: those never block longer than one quick status check and the + command survives this process restarting or losing connection, unlike a + long blocking call here which has no way to recover if interrupted. + + Do not put credentials in ``command``. Command text and output are + returned to the current step but are not persisted in the durable job + snapshot. Not every provider supports this — a batch job (e.g. + `bohr_job`) returns an error explaining that its whole command must run + at submission time instead. + """ + job = get_remote_job_status(job_id, tool_context) + if job.get("status") == "error": + return job + try: + return _service().run_job_command(job_id, command, user=user) + except Exception as exc: + current = get_remote_job_status(job_id, tool_context) + result = {"status": "error", "message": f"Remote command failed: {exc}"} + if current.get("user_control"): + result["user_control"] = current["user_control"] + return result + + +def start_remote_job_command( + job_id: str, + command: str, + tool_context: ToolContext, + user: str = "root", +) -> dict[str, Any]: + """Launch a long-running command inside a tracked interactive remote job WITHOUT blocking. + + Use this instead of `run_remote_job_command` for any real computation + (training, `vasp_std`/`mpirun`, anything that might take more than a + minute). Returns almost immediately once the command is launched in the + background; call `poll_remote_job_command` with the same `job_id` + afterward — repeatedly, across as many separate tool calls or even + separate step-executor attempts as needed — to check whether it has + finished. The command's progress is tracked durably on the job itself, so + re-attaching to this `job_id` after a step timeout, a crash, or a lost + connection always finds the same in-flight command rather than losing + track of it or risking a duplicate run. + + There is at most one in-flight background command per job; starting a + new one before polling the previous one to completion overwrites the + previous command's tracked handle. + """ + job = get_remote_job_status(job_id, tool_context) + if job.get("status") == "error": + return job + try: + return _service().start_job_command(job_id, command, user=user) + except Exception as exc: + return {"status": "error", "message": f"Failed to start remote command: {exc}"} + + +def poll_remote_job_command(job_id: str, tool_context: ToolContext) -> dict[str, Any]: + """Check on the job's most recently started background command. + + Returns `{"running": true, ...}` if it is still executing — call this + again later (e.g. after doing other work, or in a fresh step-executor + attempt after re-attaching via `get_remote_job_status`) rather than + waiting in a tight loop. Once finished, returns `{"running": false, + "exit_code": ..., "output_tail": ...}`; `output_tail` is only the last + portion of combined stdout/stderr — for the full output of a long run, + use `download_remote_job_output` on the returned `log_path`. + """ + job = get_remote_job_status(job_id, tool_context) + if job.get("status") == "error": + return job + try: + return _service().poll_job_command(job_id) + except Exception as exc: + return {"status": "error", "message": f"Failed to poll remote command: {exc}"} + + +def _resolve_workspace_child( + tool_context: ToolContext, + user_path: str, +) -> tuple[Path | None, str | None]: + """Resolve ``user_path`` against the current workspace, confining it. + + Returns ``(resolved_path, None)`` on success or ``(None, message)`` if the + workspace is unavailable or the path escapes it. Shared by upload (source) + and download (destination) so confinement logic cannot drift between them. + """ + workspace_dir = tool_context.state.get("workspace_dir") + if not workspace_dir: + return None, "No workspace_dir is available for the current step." + workspace = Path(str(workspace_dir)).resolve() + candidate = Path(user_path).expanduser() + candidate = candidate.resolve() if candidate.is_absolute() else (workspace / candidate).resolve() + if not candidate.is_relative_to(workspace): + return None, "Path must resolve inside the current workspace." + return candidate, None + + +def upload_remote_job_input( + job_id: str, + source_path: str, + destination_path: str, + tool_context: ToolContext, +) -> dict[str, Any]: + """Upload a workspace input file into a tracked interactive remote job. + + ``source_path`` must resolve inside the current workspace. Use an + absolute remote path for ``destination_path`` such as + ``/home/user/input.in``. + """ + job = get_remote_job_status(job_id, tool_context) + if job.get("status") == "error": + return job + source, error = _resolve_workspace_child(tool_context, source_path) + if error is not None: + return {"status": "error", "message": f"Upload failed: {error}"} + try: + return _service().upload_job_file(job_id, source, destination_path) + except Exception as exc: + return {"status": "error", "message": f"Upload failed: {exc}"} + + +def download_remote_job_output( + job_id: str, + source_path: str, + destination_path: str, + tool_context: ToolContext, +) -> dict[str, Any]: + """Download a file from a tracked interactive remote job into the local workspace. + + ``source_path`` is an absolute path on the remote side (e.g. + ``/home/user/CHGCAR``). ``destination_path`` must resolve inside the + current workspace. For a batch job (e.g. `bohr_job`), use + `collect_remote_job_outputs` instead once the job has succeeded. + """ + job = get_remote_job_status(job_id, tool_context) + if job.get("status") == "error": + return job + destination, error = _resolve_workspace_child(tool_context, destination_path) + if error is not None: + return {"status": "error", "message": f"Download failed: {error}"} + try: + return _service().download_job_file(job_id, source_path, destination) + except Exception as exc: + return {"status": "error", "message": f"Download failed: {exc}"} + + +def collect_remote_job_outputs( + job_id: str, + destination_path: str, + tool_context: ToolContext, +) -> dict[str, Any]: + """Pull a finished batch job's declared output files into the local workspace. + + Only valid once `get_remote_job_status` reports ``status: succeeded``. + ``destination_path`` must resolve inside the current workspace as a + directory. A repeated call after outputs are already collected is a + durable no-op that returns the same artifact list rather than + downloading twice. + """ + job = get_remote_job_status(job_id, tool_context) + if job.get("status") == "error": + return job + destination, error = _resolve_workspace_child(tool_context, destination_path) + if error is not None: + return {"status": "error", "message": f"Output collection failed: {error}"} + try: + collected = _service().collect_job_outputs(job_id, destination) + except Exception as exc: + return {"status": "error", "message": f"Output collection failed: {exc}"} + return { + "job_id": collected["job_id"], + "status": collected["status"], + "artifacts": collected.get("artifacts", []), + } diff --git a/src/matcreator/agents/execution_agent/step_executor.py b/src/matcreator/agents/execution_agent/step_executor.py index 2e822f6c..a4c67a30 100644 --- a/src/matcreator/agents/execution_agent/step_executor.py +++ b/src/matcreator/agents/execution_agent/step_executor.py @@ -18,14 +18,19 @@ from ...tools.remoteagent_tool import load_remote_a2a_agents from ...tools.util_tools import show_artifact, show_plot, show_structure from ...tools.workspace_tools import get_user_skills_root, run_bash, run_python -from .e2b_tools import ( - download_e2b_output, - get_e2b_job_status, - pause_e2b_sandbox, - run_e2b_command, +from .remote_job_tools import ( + collect_remote_job_outputs, + download_remote_job_output, + get_remote_job_status, + pause_remote_job, + poll_remote_job_command, + run_remote_job_command, + start_remote_job_command, + submit_bohr_job, + submit_bohr_sandbox, submit_e2b_sandbox, - terminate_e2b_sandbox, - upload_e2b_input, + terminate_remote_job, + upload_remote_job_input, ) logger = logging.getLogger(__name__) @@ -135,20 +140,48 @@ def _fill_missing_fields(self) -> "StepExecutorResult": 3. **If submission.json exists but outputs are missing**, reuse the same submission file (dpdispatcher is idempotent — it skips completed tasks). Do NOT regenerate submission.json. 4. **Never resubmit a job that already completed** — this wastes GPU time and creates duplicate training runs. -## Re-attaching to an existing E2B job (CRITICAL) -If your `prior_context` contains "REMOTE JOB ALREADY SUBMITTED", a tracked sandbox job -for this exact step is already running: -1. Call `get_e2b_job_status` with the given `job_id` FIRST. -2. NEVER call `submit_e2b_sandbox` for that step — it would duplicate a running job. -3. If the job finished, collect its outputs with `download_e2b_output` and report success. -4. If it is still running, call `submit_step_result(status="needs_replanning", ...)` stating +## Re-attaching to an existing remote job (CRITICAL) +If your `prior_context` contains "REMOTE JOB ALREADY SUBMITTED", a tracked remote job +(sandbox or batch job) for this exact step is already running: +1. Call `get_remote_job_status` with the given `job_id` FIRST. +2. NEVER call `submit_e2b_sandbox`, `submit_bohr_sandbox`, or `submit_bohr_job` for that + step — it would duplicate a running job. +3. If its `snapshot` contains `background_command`, a command is (or was) running in the + background — call `poll_remote_job_command` FIRST rather than issuing a new command. Never + call `start_remote_job_command`/`run_remote_job_command` again for the same computation + just because you lost track of it; re-running a non-idempotent command (e.g. a training + run) can corrupt output or double-charge compute. +4. If the job finished (`status: succeeded`), collect its outputs — `download_remote_job_output` + for an interactive sandbox, or `collect_remote_job_outputs` for a batch job — and report success. +5. If it is still running, call `submit_step_result(status="needs_replanning", ...)` stating that the job has not finished yet and quoting its job_id and status. -## User controls for E2B sandboxes -`get_e2b_job_status` may return `user_control` when the user paused or terminated -the sandbox from the UI. This does not cancel your executor. Treat it as the -user's explicit instruction: do not retry the interrupted sandbox command or -submit a replacement sandbox. Report the pause or termination accurately with +## Choosing a remote-job submit tool +- `submit_e2b_sandbox`: interactive E2B sandbox (default choice when the E2B SDK/API key + is configured). +- `submit_bohr_sandbox`: interactive sandbox via the `bohr` CLI — use only when the E2B + path is unavailable; both reach the same Bohrium sandbox platform. +- `submit_bohr_job`: batch/HPC-style job via `bohr job submit`. There is no interactive + command execution for this provider — express the entire computation in `command`, then + poll `get_remote_job_status` until `succeeded` and call `collect_remote_job_outputs`. + +## Running commands inside a sandbox: blocking vs. background (CRITICAL) +- `run_remote_job_command` BLOCKS until the command finishes, with no timeout. Only use it + for short commands expected to finish in well under a minute (`mkdir`, `grep`, checking a + file, listing a directory). +- For any real computation (a training run, `vasp_std`/`mpirun`, anything that might run + longer than a minute), use `start_remote_job_command` instead — it launches the command in + the background and returns almost immediately — then call `poll_remote_job_command` to check + on it. Do useful work in between polls (e.g. `submit_step_result(status="needs_replanning", ...)` + reporting the job is still running, rather than looping tool calls back-to-back) so a single + step doesn't sit blocked. Re-attaching to the same `job_id` later always finds the same + tracked command via `poll_remote_job_command`, so it's always safe even after a step timeout. + +## User controls for remote jobs +`get_remote_job_status` may return `user_control` when the user paused or terminated +the job from the UI. This does not cancel your executor. Treat it as the +user's explicit instruction: do not retry the interrupted command or +submit a replacement job. Report the pause or termination accurately with `submit_step_result(status="needs_replanning", replan_reason=...)`. ## MANDATORY: Always call submit_step_result @@ -264,12 +297,17 @@ def build_step_executor_agent(llm_card: LLMCard) -> LlmAgent: FunctionTool(run_python), FunctionTool(run_bash), FunctionTool(submit_e2b_sandbox), - FunctionTool(get_e2b_job_status), - FunctionTool(run_e2b_command), - FunctionTool(upload_e2b_input), - FunctionTool(download_e2b_output), - FunctionTool(pause_e2b_sandbox), - FunctionTool(terminate_e2b_sandbox), + FunctionTool(submit_bohr_sandbox), + FunctionTool(submit_bohr_job), + FunctionTool(get_remote_job_status), + FunctionTool(run_remote_job_command), + FunctionTool(start_remote_job_command), + FunctionTool(poll_remote_job_command), + FunctionTool(upload_remote_job_input), + FunctionTool(download_remote_job_output), + FunctionTool(collect_remote_job_outputs), + FunctionTool(pause_remote_job), + FunctionTool(terminate_remote_job), ALL_SKILLS_TOOLSET, FunctionTool(show_plot), FunctionTool(show_structure), diff --git a/src/matcreator/agents/execution_agent/step_executor_runner.py b/src/matcreator/agents/execution_agent/step_executor_runner.py index cd79b3f6..763c8233 100644 --- a/src/matcreator/agents/execution_agent/step_executor_runner.py +++ b/src/matcreator/agents/execution_agent/step_executor_runner.py @@ -389,13 +389,20 @@ def _remote_job_prior_context(tool_context: ToolContext, node_id: Optional[str]) remote_job = node.get("remote_job") if not isinstance(remote_job, dict) or not remote_job.get("job_id"): return None + background_note = ( + " A background command may still be in flight — call poll_remote_job_command " + "before starting a new one." + if remote_job.get("has_background_command") + else "" + ) return ( "REMOTE JOB ALREADY SUBMITTED for this step: " f"job_id={remote_job['job_id']} provider={remote_job.get('provider')} " - f"sandbox_id={remote_job.get('external_id')} last_known_status={remote_job.get('status')}. " - "Call get_e2b_job_status with this job_id to re-attach. Do NOT call submit_e2b_sandbox " - "again for this step — that job is still tracked and must not be duplicated. If it is " - "still running, report needs_replanning explaining that the job has not finished yet." + f"external_id={remote_job.get('external_id')} last_known_status={remote_job.get('status')}. " + "Call get_remote_job_status with this job_id to re-attach. Do NOT call submit_e2b_sandbox, " + "submit_bohr_sandbox, or submit_bohr_job again for this step — that job is still tracked " + "and must not be duplicated. If it is still running, report needs_replanning explaining " + f"that the job has not finished yet.{background_note}" ) diff --git a/src/matcreator/agents/thinking_agent/agent.py b/src/matcreator/agents/thinking_agent/agent.py index ac5a5e53..68de5592 100644 --- a/src/matcreator/agents/thinking_agent/agent.py +++ b/src/matcreator/agents/thinking_agent/agent.py @@ -273,8 +273,8 @@ async def run_flash_step( its timeout while the step's `remote_job` is still running on the provider. Report the job identity to the user and stop. Do NOT re-run the step to "restart" the job. To check on it later, call `run_flash_step` again with the SAME `label` and an action - that says to call `get_e2b_job_status` with that job_id and collect results if finished — - never to submit a new sandbox. + that says to call `get_remote_job_status` with that job_id and collect results if finished — + never to submit a new job. """ _NORMAL_INSTRUCTION = """ diff --git a/src/matcreator/control_plane/providers/__init__.py b/src/matcreator/control_plane/providers/__init__.py new file mode 100644 index 00000000..41793365 --- /dev/null +++ b/src/matcreator/control_plane/providers/__init__.py @@ -0,0 +1,48 @@ +"""Built-in remote-job provider adapters. + +Importing this package registers every built-in adapter. Adding a new +provider means adding one adapter module and one ``register_adapter`` call +below — nothing else in the control plane needs to change. +""" +from __future__ import annotations + +from .base import CapabilityError, RemoteJobAdapter, RemoteJobCapability, RemoteJobStatus +from .registry import get_adapter, register_adapter, registered_providers, reset_registry + + +def _register_builtin_adapters() -> None: + # Registered as lazy factories (not imported eagerly) so importing this + # package never pays for an adapter's own imports (SDK modules, env-var + # setup) until `get_adapter(...)` actually constructs one. + def _e2b_factory(): + from .e2b import E2BSandboxAdapter + + return E2BSandboxAdapter() + + def _bohr_sandbox_factory(): + from .bohr_sandbox import BohrSandboxAdapter + + return BohrSandboxAdapter() + + def _bohr_job_factory(): + from .bohr_job import BohrJobAdapter + + return BohrJobAdapter() + + register_adapter("e2b", _e2b_factory) + register_adapter("bohr_sandbox", _bohr_sandbox_factory) + register_adapter("bohr_job", _bohr_job_factory) + + +_register_builtin_adapters() + +__all__ = [ + "RemoteJobAdapter", + "RemoteJobCapability", + "RemoteJobStatus", + "CapabilityError", + "get_adapter", + "register_adapter", + "registered_providers", + "reset_registry", +] diff --git a/src/matcreator/control_plane/providers/_bohr_cli.py b/src/matcreator/control_plane/providers/_bohr_cli.py new file mode 100644 index 00000000..f9c1c369 --- /dev/null +++ b/src/matcreator/control_plane/providers/_bohr_cli.py @@ -0,0 +1,75 @@ +"""Shared subprocess boundary for `bohr`-CLI-backed provider adapters. + +Both ``bohr_sandbox`` (interactive) and ``bohr_job`` (batch) adapters shell +out to the same ``bohr`` binary and expect the same JSON envelope +(``{"ok": bool, "data": ..., "error": {...}}``), so the invocation and +error-handling logic lives here once instead of being duplicated per adapter. +""" +from __future__ import annotations + +import json +import os +import shutil +import subprocess +from typing import Any + + +class BohrCLIError(RuntimeError): + """Raised when a `bohr` CLI invocation fails or returns unusable output.""" + + +def resolve_bohr_binary() -> str: + """Resolve the `bohr` executable, honoring an explicit override.""" + return os.environ.get("BOHR_CLI_PATH") or shutil.which("bohr") or "bohr" + + +def run_bohr_json(args: list[str], *, timeout: float | None = 120) -> Any: + """Run one `bohr` CLI invocation and return its parsed ``data`` payload. + + Every invocation appends ``-o json --no-interactive -y`` so output is + machine-parseable and no command blocks on an interactive confirmation + prompt. Raises :class:`BohrCLIError` with the CLI's own error message on + failure, so callers never need to parse stderr or exit codes themselves. + """ + command = [resolve_bohr_binary(), *args, "-o", "json", "--no-interactive", "-y"] + try: + completed = subprocess.run( + command, + capture_output=True, + text=True, + timeout=timeout, + check=False, + ) + except FileNotFoundError as exc: + raise BohrCLIError("The 'bohr' CLI is not installed or not on PATH") from exc + except subprocess.TimeoutExpired as exc: + raise BohrCLIError(f"bohr {' '.join(args)} timed out after {timeout}s") from exc + + stdout = (completed.stdout or "").strip() + if not stdout: + message = (completed.stderr or "").strip() + raise BohrCLIError( + message or f"bohr {' '.join(args)} produced no output (exit {completed.returncode})" + ) + try: + payload = json.loads(stdout) + except json.JSONDecodeError as exc: + raise BohrCLIError( + f"bohr {' '.join(args)} returned non-JSON output: {stdout[:500]}" + ) from exc + + if not isinstance(payload, dict) or not payload.get("ok", False): + error = (payload or {}).get("error") if isinstance(payload, dict) else None + message = (error or {}).get("message") if isinstance(error, dict) else None + raise BohrCLIError(message or f"bohr {' '.join(args)} failed") + return payload.get("data") + + +def extract_id(data: Any, keys: tuple[str, ...]) -> str | None: + """Return the first present, truthy value among ``keys`` in a dict payload.""" + if isinstance(data, dict): + for key in keys: + value = data.get(key) + if value: + return str(value) + return None diff --git a/src/matcreator/control_plane/providers/base.py b/src/matcreator/control_plane/providers/base.py new file mode 100644 index 00000000..45edd487 --- /dev/null +++ b/src/matcreator/control_plane/providers/base.py @@ -0,0 +1,126 @@ +"""Provider-neutral adapter protocol for the remote-job control plane. + +``RemoteJobStore`` and ``RemoteJobService`` are already provider-neutral (see +``remote_jobs.py``): the persisted schema, lifecycle state machine, and +recovery bookkeeping all key off a generic ``provider`` string. This module +defines the boundary a *new* provider must implement so that +``RemoteJobService`` never needs provider-specific branches — adding a +provider means adding one adapter module plus one registration call (see +``registry.py``), nothing else in the control plane changes. +""" +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from enum import Enum +from pathlib import Path +from typing import Any + + +class RemoteJobCapability(Enum): + """Optional operations a provider adapter may support. + + ``create``/``status``/``cancel`` are mandatory for every adapter (they are + abstract methods on :class:`RemoteJobAdapter`). Everything else is gated + by a capability flag so :class:`RemoteJobService` can reject an + unsupported operation with a clear, provider-attributed error instead of + an ``AttributeError`` surfacing from deep inside an adapter. + """ + + PAUSE = "pause" + RESUME = "resume" + INTERACTIVE_EXEC = "interactive_exec" + FILE_TRANSFER = "file_transfer" + BATCH_COLLECT = "batch_collect" + + +class CapabilityError(NotImplementedError): + """Raised when a requested operation is not supported by a provider.""" + + def __init__(self, provider: str, capability: RemoteJobCapability) -> None: + super().__init__(f"Provider '{provider}' does not support '{capability.value}'") + self.provider = provider + self.capability = capability + + +@dataclass(frozen=True) +class RemoteJobStatus: + """Result of probing one external job. + + ``normalized_status`` must be one of the canonical statuses defined in + ``remote_jobs.py`` (for example ``running``, ``succeeded``, ``failed``, + ``cancelled``, ``lost``) when the provider can report a lifecycle + observation, or ``None`` when the provider can only confirm liveness + without knowing whether that changes the normalized lifecycle (e.g. an + interactive sandbox that stays "running" until an agent explicitly ends + it). ``RemoteJobService`` only transitions the durable record's status + when ``normalized_status`` is not ``None`` and differs from the job's + current status; otherwise it merges ``snapshot`` as a non-lifecycle + observation. + """ + + normalized_status: str | None + snapshot: dict[str, Any] = field(default_factory=dict) + error: str | None = None + + +class RemoteJobAdapter(ABC): + """Boundary between the control plane and one external job provider. + + Subclasses implement only what their provider can actually do. Declaring + ``capabilities`` tells :class:`RemoteJobService` which of the optional + methods below are safe to call; the default implementations below raise + :class:`CapabilityError` as a safety net for a capability that was + declared but never overridden. + """ + + provider: str + capabilities: frozenset[RemoteJobCapability] = frozenset() + # Interval RemoteJobMonitor should wait between reconciliations of jobs + # owned by this provider. Batch/HPC-style providers whose status only + # changes on the order of minutes can declare a much longer interval than + # an interactively addressable sandbox. + poll_interval_seconds: float = 15.0 + + @abstractmethod + def create(self, spec: dict[str, Any]) -> str: + """Create one external job/sandbox and return its provider-side ID.""" + + @abstractmethod + def status(self, external_id: str) -> RemoteJobStatus: + """Probe one external job for liveness and/or lifecycle status.""" + + @abstractmethod + def cancel(self, external_id: str) -> None: + """Irreversibly stop/delete one external job.""" + + def pause(self, external_id: str) -> None: + raise CapabilityError(self.provider, RemoteJobCapability.PAUSE) + + def resume(self, external_id: str) -> None: + raise CapabilityError(self.provider, RemoteJobCapability.RESUME) + + def run_command(self, external_id: str, command: str, *, user: str = "root") -> dict[str, Any]: + raise CapabilityError(self.provider, RemoteJobCapability.INTERACTIVE_EXEC) + + def upload_file(self, external_id: str, source: str | Path, destination: str) -> None: + raise CapabilityError(self.provider, RemoteJobCapability.FILE_TRANSFER) + + def download_file( + self, + external_id: str, + source: str, + destination: str | Path, + *, + user: str | None = None, + ) -> Path: + raise CapabilityError(self.provider, RemoteJobCapability.FILE_TRANSFER) + + def collect_outputs(self, external_id: str, destination_dir: str | Path) -> list[dict[str, Any]]: + """Pull a finished batch job's declared output files to ``destination_dir``. + + Returns a list of ``{"source": ..., "destination": ...}`` records + describing what was collected, so the caller can persist them as job + artifacts. + """ + raise CapabilityError(self.provider, RemoteJobCapability.BATCH_COLLECT) diff --git a/src/matcreator/control_plane/providers/bohr_job.py b/src/matcreator/control_plane/providers/bohr_job.py new file mode 100644 index 00000000..ddc64822 --- /dev/null +++ b/src/matcreator/control_plane/providers/bohr_job.py @@ -0,0 +1,109 @@ +"""Batch/HPC-style adapter over `bohr job` (submit/describe/download/terminate). + +Inputs are staged once at submission time (``--input_directory``); there is +no interactive exec or incremental file transfer, matching how HPC batch +schedulers work (submit, poll a queue, collect outputs once terminal). Only +single-job submission is supported here — `bohr job_group` fan-out (multiple +jobs sharing one group) is a possible future adapter, not this one. +""" +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from ._bohr_cli import BohrCLIError, extract_id, run_bohr_json +from .base import RemoteJobAdapter, RemoteJobCapability, RemoteJobStatus + +# `bohr job list`/`describe` report a lowercase `phase` string. Map it onto +# the canonical statuses defined in remote_jobs.py. Any phase not listed here +# (e.g. a future platform phase) leaves normalized_status as None so the +# service records an observation instead of guessing a lifecycle transition. +_PHASE_TO_NORMALIZED = { + "pending": "queued", + "scheduling": "queued", + "running": "running", + "completed": "succeeded", + "failed": "failed", + # "stopped" is the phase used for a job terminated by the user (see + # `bohr job terminate`); "cancelled" is the closest existing canonical + # status for a non-failure, user-initiated stop. + "stopped": "cancelled", +} + +_REQUIRED_SPEC_FIELDS = ("project_id", "job_name", "machine_type", "image_address", "command") + + +class BohrJobAdapter(RemoteJobAdapter): + provider = "bohr_job" + capabilities = frozenset({RemoteJobCapability.BATCH_COLLECT}) + # Batch job status changes on the order of minutes, not seconds; poll far + # less often than an interactive sandbox to avoid hammering the platform. + poll_interval_seconds = 60.0 + + def create(self, spec: dict[str, Any]) -> str: + missing = [name for name in _REQUIRED_SPEC_FIELDS if not spec.get(name)] + if missing: + raise ValueError(f"bohr_job spec is missing required field(s): {', '.join(missing)}") + args = [ + "job", + "submit", + "--project_id", + str(spec["project_id"]), + "--job_name", + str(spec["job_name"]), + "--machine_type", + str(spec["machine_type"]), + "--image_address", + str(spec["image_address"]), + "--command", + str(spec["command"]), + ] + if spec.get("input_directory"): + args += ["--input_directory", str(spec["input_directory"])] + if spec.get("log_file"): + args += ["--log_file", str(spec["log_file"])] + if spec.get("result_path"): + args += ["--result_path", str(spec["result_path"])] + if spec.get("max_run_time"): + args += ["--max_run_time", str(spec["max_run_time"])] + if spec.get("nnode"): + args += ["--nnode", str(spec["nnode"])] + if spec.get("max_reschedule_times") is not None: + args += ["--max_reschedule_times", str(spec["max_reschedule_times"])] + if spec.get("job_group_id"): + args += ["--job_group_id", str(spec["job_group_id"])] + + data = run_bohr_json(args) + bohr_id = extract_id(data, ("bohrId", "bohr_id", "bohrID", "id", "jobId", "jobID")) + if not bohr_id: + keys = sorted(data) if isinstance(data, dict) else type(data).__name__ + raise BohrCLIError( + f"bohr job submit did not return a Bohr job ID (data: {keys})" + ) + return bohr_id + + def status(self, external_id: str) -> RemoteJobStatus: + data = run_bohr_json(["job", "describe", "-i", str(external_id)]) or {} + phase = str(data.get("phase", "")).lower() + normalized = _PHASE_TO_NORMALIZED.get(phase) + error = None + if phase == "failed": + error = str(data.get("errorInfo") or "").strip() or None + return RemoteJobStatus( + normalized_status=normalized, + snapshot={"phase": phase or None, "terminal": bool(data.get("terminal", False))}, + error=error, + ) + + def cancel(self, external_id: str) -> None: + run_bohr_json(["job", "terminate", "--id", str(external_id), "--no-wait"]) + + def collect_outputs(self, external_id: str, destination_dir: str | Path) -> list[dict[str, Any]]: + dest = Path(destination_dir).expanduser().resolve() + dest.mkdir(parents=True, exist_ok=True) + run_bohr_json(["job", "download", "-i", str(external_id), "--out", str(dest)]) + return [ + {"source": external_id, "destination": str(path)} + for path in sorted(dest.rglob("*")) + if path.is_file() + ] diff --git a/src/matcreator/control_plane/providers/bohr_sandbox.py b/src/matcreator/control_plane/providers/bohr_sandbox.py new file mode 100644 index 00000000..b3dc7bbe --- /dev/null +++ b/src/matcreator/control_plane/providers/bohr_sandbox.py @@ -0,0 +1,118 @@ +"""Interactive adapter over `bohr sandbox` (create/exec/files/describe/delete). + +Mirrors the E2B adapter's capability surface (interactive exec + file +transfer) so agent tools can treat a Bohrium CLI sandbox the same way as an +E2B one. The installed `bohr` CLI has no sandbox pause/resume subcommand +(only create/delete/describe/exec/files/list/...), so +``RemoteJobCapability.PAUSE`` is intentionally not declared here — add it if +a future CLI version exposes one. +""" +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from ._bohr_cli import BohrCLIError, extract_id, run_bohr_json +from .base import RemoteJobAdapter, RemoteJobCapability, RemoteJobStatus + +# Sandboxes reporting one of these as their describe-output status/state are +# no longer usable; treat them the same as an E2B "unreachable" liveness +# probe failure so the monitor marks the job lost rather than looping. +_UNREACHABLE_STATUSES = {"deleted", "terminated", "stopped", "killed"} + + +class BohrSandboxAdapter(RemoteJobAdapter): + provider = "bohr_sandbox" + capabilities = frozenset({RemoteJobCapability.INTERACTIVE_EXEC, RemoteJobCapability.FILE_TRANSFER}) + poll_interval_seconds = 15.0 + + def create(self, spec: dict[str, Any]) -> str: + project_id = spec.get("project_id") + if not project_id: + raise ValueError("bohr_sandbox spec requires 'project_id'") + template = spec.get("template") + if not template: + # Never fall back to the CLI's default template (sdbxagent): it + # silently creates the wrong (and possibly costlier) sandbox. + raise ValueError("bohr_sandbox spec requires 'template'") + args = ["sandbox", "create", "--template", str(template), "--project-id", str(project_id)] + if spec.get("timeout"): + args += ["--timeout", str(spec["timeout"])] + if spec.get("image"): + args += ["--image", str(spec["image"])] + if spec.get("gpu"): + args += ["--gpu", str(spec["gpu"])] + if spec.get("never_timeout"): + args.append("--never-timeout") + if spec.get("mount_user_storage"): + args.append("--mount-user-storage") + if spec.get("share_subpath"): + args += ["--share-subpath", str(spec["share_subpath"])] + if spec.get("session_id"): + args += ["--session-id", str(spec["session_id"])] + for key, value in dict(spec.get("env") or {}).items(): + args += ["--env", f"{key}={value}"] + + data = run_bohr_json(args) + # The CLI is inconsistent about the ID key across subcommands: + # `create`/`exec` return "sandboxID", `files read` returns + # "sandbox_id" — accept every observed spelling. + sandbox_id = extract_id(data, ("sandbox_id", "sandboxID", "sandboxId", "id")) + if not sandbox_id: + keys = sorted(data) if isinstance(data, dict) else type(data).__name__ + raise BohrCLIError( + f"bohr sandbox create did not return a sandbox ID (data: {keys})" + ) + return sandbox_id + + def status(self, external_id: str) -> RemoteJobStatus: + data = run_bohr_json(["sandbox", "describe", external_id]) or {} + raw_status = str(data.get("status") or data.get("state") or "").lower() + normalized = "lost" if raw_status in _UNREACHABLE_STATUSES else None + return RemoteJobStatus( + normalized_status=normalized, + snapshot={ + "provider_status": "unreachable" if normalized else "reachable", + "raw_status": raw_status or None, + }, + error=None, + ) + + def cancel(self, external_id: str) -> None: + run_bohr_json(["sandbox", "delete", external_id, "--force"]) + + def run_command(self, external_id: str, command: str, *, user: str = "root") -> dict[str, Any]: + # `bohr sandbox exec` caps a command at 90s by default; pass + # `--timeout 0` to disable that CLI-side cap, and also don't bound our + # own subprocess wait, so a long-running command isn't silently + # truncated. This matches the E2B adapter's `timeout=0` semantics + # (see e2b.py run_command) so command duration behaves the same + # regardless of which interactive sandbox provider is in use. + data = run_bohr_json( + ["sandbox", "exec", external_id, "--command", command, "--user", user, "--timeout", "0"], + timeout=None, + ) or {} + return { + "stdout": str(data.get("stdout", "")), + "stderr": str(data.get("stderr", "")), + "exit_code": data.get("exit_code", data.get("exitCode")), + } + + def upload_file(self, external_id: str, source: str | Path, destination: str) -> None: + source_path = Path(source).expanduser().resolve() + if not source_path.is_file(): + raise FileNotFoundError(source_path) + run_bohr_json(["sandbox", "files", "write", external_id, destination, "--source", str(source_path)]) + + def download_file( + self, + external_id: str, + source: str, + destination: str | Path, + *, + user: str | None = None, + ) -> Path: + dest_path = Path(destination).expanduser().resolve() + dest_path.parent.mkdir(parents=True, exist_ok=True) + run_bohr_json(["sandbox", "files", "read", external_id, source, "--destination", str(dest_path)]) + return dest_path diff --git a/src/matcreator/control_plane/e2b.py b/src/matcreator/control_plane/providers/e2b.py similarity index 63% rename from src/matcreator/control_plane/e2b.py rename to src/matcreator/control_plane/providers/e2b.py index fc0ca055..e3243893 100644 --- a/src/matcreator/control_plane/e2b.py +++ b/src/matcreator/control_plane/providers/e2b.py @@ -6,6 +6,8 @@ from typing import Any import os +from .base import RemoteJobAdapter, RemoteJobCapability, RemoteJobStatus + # e2b SDK >=2.20 validates the API key format client-side (requires the # "e2b_" + hex pattern, see e2b.api.validate_api_key) and raises # AuthenticationException in ApiClient.__init__ before any network call. @@ -25,6 +27,27 @@ class E2BUnavailableError(RuntimeError): """Raised when the optional E2B SDK is unavailable at runtime.""" +@dataclass(frozen=True) +class E2BConnectionConfig: + """Server-side E2B/Bohrium endpoint configuration for one sandbox request.""" + + api_key: str + api_url: str + project_id: str + template: str + + def to_spec_dict(self, *, timeout: int = 600, lifecycle: dict[str, Any] | None = None) -> dict[str, Any]: + """Build the generic ``spec`` dict passed to ``E2BSandboxAdapter.create``.""" + return { + "template": self.template, + "api_key": self.api_key, + "api_url": self.api_url, + "project_id": self.project_id, + "timeout": timeout, + "lifecycle": lifecycle or {}, + } + + @dataclass(frozen=True) class E2BSandboxSpec: """Validated inputs for creating one E2B sandbox.""" @@ -54,10 +77,36 @@ def create_kwargs(self) -> dict[str, Any]: "metadata": self.metadata, } + @classmethod + def from_dict(cls, spec: dict[str, Any]) -> "E2BSandboxSpec": + """Build a validated spec from the generic dict a submit tool provides.""" + return cls( + template=str(spec.get("template", "")), + api_key=str(spec.get("api_key", "")), + api_url=str(spec.get("api_url", "")), + project_id=str(spec.get("project_id", "")), + timeout=int(spec.get("timeout", 600)), + lifecycle=dict(spec.get("lifecycle") or {}), + metadata=dict(spec.get("metadata") or {}), + ) + # The backend E2B SDK is imported lazily to avoid a hard dependency on the SDK for users who don't need it. The E2BSandboxAdapter class wraps the SDK and provides a simple interface for creating, connecting to, and managing E2B sandboxes. -class E2BSandboxAdapter: +class E2BSandboxAdapter(RemoteJobAdapter): """Small boundary around the E2B SDK with no SDK import at module load.""" + provider = "e2b" + # Resume is not implemented: the E2B SDK's "pause" produces a snapshot that + # is restored by simply connecting to the same sandbox_id again, so there + # is no separate resume operation to expose. + capabilities = frozenset( + { + RemoteJobCapability.PAUSE, + RemoteJobCapability.INTERACTIVE_EXEC, + RemoteJobCapability.FILE_TRANSFER, + } + ) + poll_interval_seconds = 15.0 + @staticmethod def _sandbox_class(): try: @@ -68,8 +117,20 @@ def _sandbox_class(): ) from exc return Sandbox - def create(self, spec: E2BSandboxSpec) -> str: - sandbox = self._sandbox_class().create(**spec.create_kwargs()) + def create(self, spec: dict[str, Any]) -> str: + sandbox_spec = E2BSandboxSpec.from_dict(spec) + try: + sandbox = self._sandbox_class().create(**sandbox_spec.create_kwargs()) + except (E2BConfigurationError, E2BUnavailableError): + raise + except Exception as exc: + # Include every non-secret request parameter so a provider-side + # error (404 wrong URL, unknown template, bad project) is + # diagnosable from the durable job record alone. + raise RuntimeError( + f"{exc} [create against api_url={sandbox_spec.api_url!r}, " + f"template={sandbox_spec.template!r}, project_id={sandbox_spec.project_id!r}]" + ) from exc sandbox_id = getattr(sandbox, "sandbox_id", "") if not sandbox_id: raise RuntimeError("E2B create returned a sandbox without sandbox_id") @@ -170,6 +231,9 @@ def pause(self, sandbox_id: str) -> None: def terminate(self, sandbox_id: str) -> None: self._connect(sandbox_id).kill() + def cancel(self, external_id: str) -> None: + self.terminate(external_id) + def probe(self, sandbox_id: str) -> dict[str, Any]: """Confirm an active sandbox is reachable without changing its files.""" result = self.run_command(sandbox_id, "true") @@ -177,7 +241,31 @@ def probe(self, sandbox_id: str) -> dict[str, Any]: raise RuntimeError(result["stderr"] or "E2B sandbox liveness probe failed") return {"provider_status": "reachable", "probe": result} + def status(self, external_id: str) -> RemoteJobStatus: + """Report liveness only: an E2B sandbox stays "running" until an agent + or user explicitly pauses/terminates it, so there is no independent + lifecycle status for the monitor to observe beyond reachability. + """ + snapshot = self.probe(external_id) + return RemoteJobStatus(normalized_status=None, snapshot=snapshot, error=None) + def _connect(self, sandbox_id: str): if not sandbox_id: raise ValueError("sandbox_id is required") - return self._sandbox_class().connect(sandbox_id) \ No newline at end of file + # Reconnects (run_command/pause/kill, possibly from a fresh process + # long after create) must target the same Bohrium endpoint as create. + # The SDK would silently fall back to the public e2b.dev API when + # called bare, so pass the configured connection explicitly and fail + # loudly when it is missing. + api_key = os.environ.get("E2B_API_KEY", "") + api_url = os.environ.get("E2B_API_URL", "") + if not api_key or not api_url: + raise E2BConfigurationError( + "E2B_API_KEY and E2B_API_URL must be set to connect to sandbox " + f"'{sandbox_id}'" + ) + opts: dict[str, Any] = {"api_key": api_key, "api_url": api_url} + project_id = os.environ.get("BOHRIUM_PROJECT_ID", "") + if project_id: + opts["headers"] = {"X-Project-Id": project_id} + return self._sandbox_class().connect(sandbox_id, **opts) \ No newline at end of file diff --git a/src/matcreator/control_plane/providers/registry.py b/src/matcreator/control_plane/providers/registry.py new file mode 100644 index 00000000..ed01ad3c --- /dev/null +++ b/src/matcreator/control_plane/providers/registry.py @@ -0,0 +1,50 @@ +"""Registry mapping remote-job provider names to adapter instances. + +Adding a new provider means adding one adapter module implementing +``RemoteJobAdapter`` and one ``register_adapter`` call in +``providers/__init__.py`` — nothing else in the control plane changes. +""" +from __future__ import annotations + +from typing import Callable + +from .base import RemoteJobAdapter + +_FACTORIES: dict[str, Callable[[], RemoteJobAdapter]] = {} +_INSTANCES: dict[str, RemoteJobAdapter] = {} + + +def register_adapter(provider: str, factory: Callable[[], RemoteJobAdapter]) -> None: + """Register a lazy factory for one provider's adapter. + + Factories are not called until first use, so importing the registry never + imports an optional provider SDK or shells out to a CLI. + """ + if not provider: + raise ValueError("provider is required") + _FACTORIES[provider] = factory + _INSTANCES.pop(provider, None) + + +def get_adapter(provider: str) -> RemoteJobAdapter: + """Return the (lazily constructed, cached) adapter for ``provider``.""" + if provider not in _FACTORIES: + raise KeyError(f"No remote-job adapter is registered for provider '{provider}'") + if provider not in _INSTANCES: + adapter = _FACTORIES[provider]() + if adapter.provider != provider: + raise ValueError( + f"Adapter registered for '{provider}' reports provider '{adapter.provider}'" + ) + _INSTANCES[provider] = adapter + return _INSTANCES[provider] + + +def registered_providers() -> list[str]: + return sorted(_FACTORIES) + + +def reset_registry() -> None: + """Test helper: drop all registrations and cached adapter instances.""" + _FACTORIES.clear() + _INSTANCES.clear() diff --git a/src/matcreator/control_plane/remote_job_monitor.py b/src/matcreator/control_plane/remote_job_monitor.py index dadb2967..e2cb5c45 100644 --- a/src/matcreator/control_plane/remote_job_monitor.py +++ b/src/matcreator/control_plane/remote_job_monitor.py @@ -10,10 +10,14 @@ class RemoteJobMonitor: - """Probe active E2B sandboxes with bounded retry backoff. + """Probe active remote jobs of any registered provider with bounded retry backoff. Job records are durable; this monitor's due times are intentionally process - local. On a restart its empty schedule reconciles every active sandbox once. + local. On a restart its empty schedule reconciles every active job once. + Each provider adapter declares its own ``poll_interval_seconds`` (see + ``providers/base.py``), so a batch/HPC-style provider can poll far less + often than an interactive sandbox with no change needed here — adding a + provider is a pure plugin. """ def __init__( @@ -45,30 +49,46 @@ async def run(self) -> None: def stop(self) -> None: self._stop.set() + def _base_interval(self, provider: str) -> float: + """Return the owning adapter's preferred poll cadence for one job. + + Resolved through ``self.service.adapter_for`` so this honors the same + adapter overrides (e.g. in tests) as every other service operation, + instead of querying the global registry directly. Falls back to this + monitor's own tick interval if the provider is unregistered (e.g. a + record left over from a removed plugin), so a missing adapter never + breaks reconciliation of other jobs. + """ + try: + return self.service.adapter_for(provider).poll_interval_seconds + except KeyError: + return self.interval_seconds + async def reconcile_once(self) -> list[dict[str, Any]]: now = time.monotonic() updates: list[dict[str, Any]] = [] active_ids: set[str] = set() - for job in self.store.list_active_jobs(provider="e2b"): + for job in self.store.list_active_jobs(): job_id = job["job_id"] active_ids.add(job_id) if job["status"] not in {"queued", "running", "submitting", "resuming"}: continue if now < self._next_due.get(job_id, 0): continue - updated = await asyncio.to_thread(self.service.reconcile_e2b, job_id) + base_interval = self._base_interval(job["provider"]) + updated = await asyncio.to_thread(self.service.reconcile_job, job_id) updates.append(updated) if updated["snapshot"].get("provider_status") == "unreachable": failures = self._failures.get(job_id, 0) + 1 self._failures[job_id] = failures - delay = min(self.interval_seconds * (2 ** (failures - 1)), self.max_backoff_seconds) + delay = min(base_interval * (2 ** (failures - 1)), self.max_backoff_seconds) else: self._failures.pop(job_id, None) - delay = self.interval_seconds + delay = base_interval self._next_due[job_id] = time.monotonic() + delay stale_ids = set(self._next_due) - active_ids for job_id in stale_ids: self._next_due.pop(job_id, None) self._failures.pop(job_id, None) - return updates \ No newline at end of file + return updates diff --git a/src/matcreator/control_plane/remote_job_service.py b/src/matcreator/control_plane/remote_job_service.py index 1e1cf060..b7162dc8 100644 --- a/src/matcreator/control_plane/remote_job_service.py +++ b/src/matcreator/control_plane/remote_job_service.py @@ -1,188 +1,389 @@ -"""Provider operations coordinated with durable remote-job records.""" +"""Provider operations coordinated with durable remote-job records. + +``RemoteJobService`` never branches on a provider name itself: every +operation looks up the adapter registered for ``job["provider"]`` (see +``providers/registry.py``) and checks its declared capabilities before +calling an optional method. Adding a new remote-job provider is therefore a +pure plugin: implement ``RemoteJobAdapter`` and register it — no changes +needed here. +""" from __future__ import annotations -from dataclasses import dataclass +import base64 +import time from pathlib import Path from typing import Any -from .e2b import E2BSandboxAdapter, E2BSandboxSpec +from .providers import CapabilityError, RemoteJobAdapter, RemoteJobCapability, get_adapter from .remote_jobs import RemoteJobStore -@dataclass(frozen=True) -class E2BConnectionConfig: - api_key: str - api_url: str - project_id: str - template: str - - class RemoteJobService: """Coordinates provider side effects with persisted job state.""" - def __init__(self, store: RemoteJobStore, *, e2b_adapter: E2BSandboxAdapter | None = None) -> None: + def __init__( + self, + store: RemoteJobStore, + *, + adapter_overrides: dict[str, RemoteJobAdapter] | None = None, + ) -> None: + """Create a service backed by ``store``. + + ``adapter_overrides`` lets callers (chiefly tests) inject a fake + adapter for one provider without mutating the global registry; + providers not present in the override map fall back to + ``providers.get_adapter``. + """ self.store = store - self.e2b_adapter = e2b_adapter or E2BSandboxAdapter() + self._adapter_overrides = dict(adapter_overrides or {}) + + def adapter_for(self, provider: str) -> RemoteJobAdapter: + """Resolve the adapter this service would use for ``provider``. + + Public so callers that need adapter metadata without performing an + operation (e.g. :class:`RemoteJobMonitor` reading + ``poll_interval_seconds``) resolve through the same override-aware + lookup as every service method, instead of querying the global + registry directly and silently ignoring test overrides. + """ + if provider in self._adapter_overrides: + return self._adapter_overrides[provider] + return get_adapter(provider) + + def _adapter(self, provider: str) -> RemoteJobAdapter: + return self.adapter_for(provider) - def submit_e2b( + def _get_job(self, job_id: str) -> dict[str, Any]: + job = self.store.get_job(job_id) + if job is None: + raise KeyError(f"Remote job '{job_id}' was not found") + if not job["external_id"]: + raise ValueError(f"Remote job '{job_id}' has no provider-side ID") + return job + + def submit_job( self, *, owner_id: str, session_id: str, + provider: str, idempotency_key: str, - connection: E2BConnectionConfig, + spec: dict[str, Any], node_id: str | None = None, step_number: int | None = None, - timeout: int = 600, - lifecycle: dict[str, Any] | None = None, - metadata: dict[str, str] | None = None, output_dir: str | None = None, + persisted_specification: dict[str, Any] | None = None, ) -> dict[str, Any]: - """Create an E2B sandbox once and persist the resulting sandbox ID. + """Create one external job/sandbox once and persist its external ID. - The stored specification intentionally excludes ``api_key``. Replays with - the same idempotency key return the already-created record instead of - creating another sandbox. + ``spec`` is passed to the provider adapter's ``create`` verbatim and + may contain secrets (e.g. an API key). ``persisted_specification`` — + which must never contain secrets — is stored in the durable record + instead; if omitted, ``spec`` itself is persisted, so callers whose + spec has no secrets can rely on the default. Replays with the same + idempotency key return the already-created record instead of + creating a second external job; a record that failed before ever + acquiring an external ID is reset and retried instead. """ job = self.store.create_job( owner_id=owner_id, session_id=session_id, - provider="e2b", + provider=provider, idempotency_key=idempotency_key, node_id=node_id, step_number=step_number, - specification={ - "template": connection.template, - "api_url": connection.api_url, - "project_id": connection.project_id, - "timeout": timeout, - "lifecycle": lifecycle or {}, - "metadata": metadata or {}, - }, + specification=persisted_specification if persisted_specification is not None else spec, output_dir=output_dir, ) + if job["status"] == "failed" and not job["external_id"]: + # The previous attempt died before the provider handed back an + # external ID, so nothing external exists to duplicate — retry + # instead of returning the poisoned record forever. + job = self.store.reset_failed_job_for_retry(job["job_id"]) if job["external_id"] or job["status"] != "created": return job + adapter = self._adapter(provider) submitting = self.store.transition_job(job["job_id"], "submitting") try: - sandbox_id = self.e2b_adapter.create( - E2BSandboxSpec( - template=connection.template, - api_key=connection.api_key, - api_url=connection.api_url, - project_id=connection.project_id, - timeout=timeout, - lifecycle=lifecycle or {}, - metadata=metadata or {}, - ) - ) + external_id = adapter.create(spec) except Exception as exc: return self.store.transition_job( job["job_id"], "failed", - error=f"E2B sandbox creation failed: {exc}", + error=f"{provider} job creation failed: {exc}", expected_revision=submitting["state_revision"], ) + + # Give the adapter a chance to report an initial lifecycle status + # (e.g. a batch provider that queues before it runs) instead of + # always assuming "running". + initial_status = "running" + initial_snapshot: dict[str, Any] = {"provider_status": "running"} + try: + probe = adapter.status(external_id) + except Exception: + probe = None + if probe is not None: + if probe.normalized_status: + initial_status = probe.normalized_status + initial_snapshot = {**initial_snapshot, **probe.snapshot} + return self.store.transition_job( job["job_id"], - "running", - external_id=sandbox_id, - snapshot={"provider_status": "running", "sandbox_id": sandbox_id}, + initial_status, + external_id=external_id, + snapshot=initial_snapshot, expected_revision=submitting["state_revision"], ) - def pause_e2b(self, job_id: str) -> dict[str, Any]: - self._get_e2b_job(job_id) + def pause_job(self, job_id: str) -> dict[str, Any]: + job = self._get_job(job_id) + adapter = self._adapter(job["provider"]) + if RemoteJobCapability.PAUSE not in adapter.capabilities: + raise CapabilityError(job["provider"], RemoteJobCapability.PAUSE) requested = self.store.transition_job(job_id, "pause_requested") try: - self.e2b_adapter.pause(requested["external_id"]) + adapter.pause(requested["external_id"]) except Exception as exc: return self.store.transition_job( job_id, "failed", - error=f"E2B pause failed: {exc}", + error=f"{job['provider']} pause failed: {exc}", expected_revision=requested["state_revision"], ) - return self.store.transition_job( - job_id, "paused", expected_revision=requested["state_revision"] - ) + return self.store.transition_job(job_id, "paused", expected_revision=requested["state_revision"]) - def terminate_e2b(self, job_id: str) -> dict[str, Any]: - self._get_e2b_job(job_id) + def resume_job(self, job_id: str) -> dict[str, Any]: + job = self._get_job(job_id) + adapter = self._adapter(job["provider"]) + if RemoteJobCapability.RESUME not in adapter.capabilities: + raise CapabilityError(job["provider"], RemoteJobCapability.RESUME) + requested = self.store.transition_job(job_id, "resume_requested") + resuming = self.store.transition_job(job_id, "resuming", expected_revision=requested["state_revision"]) + try: + adapter.resume(resuming["external_id"]) + except Exception as exc: + return self.store.transition_job( + job_id, + "failed", + error=f"{job['provider']} resume failed: {exc}", + expected_revision=resuming["state_revision"], + ) + return self.store.transition_job(job_id, "running", expected_revision=resuming["state_revision"]) + + def terminate_job(self, job_id: str) -> dict[str, Any]: + job = self._get_job(job_id) + adapter = self._adapter(job["provider"]) requested = self.store.transition_job(job_id, "terminate_requested") try: - self.e2b_adapter.terminate(requested["external_id"]) + adapter.cancel(requested["external_id"]) except Exception as exc: return self.store.transition_job( job_id, "lost", - error=f"E2B termination could not be confirmed: {exc}", + error=f"{job['provider']} termination could not be confirmed: {exc}", expected_revision=requested["state_revision"], ) - return self.store.transition_job( - job_id, "terminated", expected_revision=requested["state_revision"] - ) + return self.store.transition_job(job_id, "terminated", expected_revision=requested["state_revision"]) - def pause_active_session_e2b_jobs(self, *, owner_id: str, session_id: str) -> list[dict[str, Any]]: - """Request a provider pause for each active E2B job in one session.""" + def pause_active_session_jobs(self, *, owner_id: str, session_id: str) -> list[dict[str, Any]]: + """Request a provider pause for each active, pausable job in one session.""" results: list[dict[str, Any]] = [] for job in self.store.list_jobs(owner_id=owner_id, session_id=session_id): - if job["provider"] != "e2b" or job["status"] not in {"queued", "running"}: + if job["status"] not in {"queued", "running"}: continue try: - results.append(self.pause_e2b(job["job_id"])) + adapter = self._adapter(job["provider"]) + except KeyError: + continue + if RemoteJobCapability.PAUSE not in adapter.capabilities: + continue + try: + results.append(self.pause_job(job["job_id"])) except Exception as exc: - results.append( - { - "job_id": job["job_id"], - "status": job["status"], - "pause_error": str(exc), - } - ) + results.append({"job_id": job["job_id"], "status": job["status"], "pause_error": str(exc)}) return results - def reconcile_e2b(self, job_id: str) -> dict[str, Any]: - """Probe a persisted active sandbox after a process restart or refresh.""" - job = self._get_e2b_job(job_id) + def reconcile_job(self, job_id: str) -> dict[str, Any]: + """Probe a persisted active job after a process restart or refresh. + + Transitions the durable status only when the adapter reports a + normalized status that differs from the current one; otherwise the + probe result is merged as a non-lifecycle observation. An illegal + transition reported by a confused/stale adapter observation falls + back to an observation rather than raising, since a monitor loop + must never crash on one bad probe. + """ + job = self._get_job(job_id) if job["status"] not in {"queued", "running", "submitting", "resuming"}: return job + adapter = self._adapter(job["provider"]) try: - snapshot = self.e2b_adapter.probe(job["external_id"]) + probe = adapter.status(job["external_id"]) except Exception as exc: return self.store.record_observation( job_id, snapshot={"provider_status": "unreachable"}, - error=f"E2B reconciliation failed: {exc}", + error=f"{job['provider']} reconciliation failed: {exc}", expected_revision=job["state_revision"], ) + if probe.normalized_status and probe.normalized_status != job["status"]: + try: + return self.store.transition_job( + job_id, + probe.normalized_status, + snapshot=probe.snapshot, + error=probe.error, + expected_revision=job["state_revision"], + ) + except ValueError: + # Provider reported a status this job's current state cannot + # legally move to (e.g. a stale/out-of-order observation). + # Recording it as telemetry is always safe; only a lifecycle + # transition needs the strict check. + pass return self.store.record_observation( job_id, - snapshot=snapshot, - error=None, + snapshot=probe.snapshot, + error=probe.error, expected_revision=job["state_revision"], ) - def run_e2b_command(self, job_id: str, command: str, *, user: str = "root") -> dict[str, Any]: - """Run one command inside a tracked E2B sandbox without persisting command text.""" - job = self._get_e2b_job(job_id) + def run_job_command(self, job_id: str, command: str, *, user: str = "root") -> dict[str, Any]: + """Run one command inside a tracked interactive job without persisting command text. + + This blocks the caller for the command's full duration with no + timeout of its own. Fine for short commands (seconds); for anything + that might run more than a minute or two, use + ``start_job_command``/``poll_job_command`` instead, which never + blocks longer than one bounded status check and durably survives a + process restart mid-command. + """ + job = self._get_job(job_id) if job["status"] not in {"queued", "running", "resuming"}: - raise ValueError(f"E2B job '{job_id}' cannot run commands while {job['status']}") - result = self.e2b_adapter.run_command(job["external_id"], command, user=user) + raise ValueError(f"Job '{job_id}' cannot run commands while {job['status']}") + adapter = self._adapter(job["provider"]) + if RemoteJobCapability.INTERACTIVE_EXEC not in adapter.capabilities: + raise CapabilityError(job["provider"], RemoteJobCapability.INTERACTIVE_EXEC) + result = adapter.run_command(job["external_id"], command, user=user) self.store.merge_observation( job_id, - snapshot={"provider_status": "reachable", "last_command_exit_code": result["exit_code"]}, + snapshot={"provider_status": "reachable", "last_command_exit_code": result.get("exit_code")}, error=None, ) return result - def upload_e2b_file(self, job_id: str, source: str | Path, destination: str) -> dict[str, Any]: - """Upload one local input file into a tracked E2B sandbox.""" - job = self._get_e2b_job(job_id) + def start_job_command(self, job_id: str, command: str, *, user: str = "root") -> dict[str, Any]: + """Launch one command in the background inside a tracked interactive job. + + Built entirely on the existing ``run_command`` capability — no new + adapter method or capability is required, so this works for any + current or future ``INTERACTIVE_EXEC`` provider automatically. The + launch call itself returns almost immediately (only the wrapper + shell backgrounds and detaches; it does not wait for ``command`` to + finish). ``command`` is base64-encoded before being embedded in the + wrapper so arbitrary shell content (quotes, `$`, backticks, newlines) + can never break out of or reinterpret the wrapper script. + + The command's stdout/stderr and exit code are redirected to marker + files whose paths are derived only from ``job_id`` and persisted in + the job's durable snapshot. This is what makes the command + recoverable: if this process crashes or the connection drops while + the command is still running, a fresh process re-attaches to the + same job, reads the same marker-file paths from the durable record, + and calls ``poll_job_command`` — it never has to guess whether an + earlier command already ran or re-issue it, which would be unsafe + for a non-idempotent computation. + """ + job = self._get_job(job_id) + if job["status"] not in {"queued", "running", "resuming"}: + raise ValueError(f"Job '{job_id}' cannot run commands while {job['status']}") + adapter = self._adapter(job["provider"]) + if RemoteJobCapability.INTERACTIVE_EXEC not in adapter.capabilities: + raise CapabilityError(job["provider"], RemoteJobCapability.INTERACTIVE_EXEC) + + marker = f"/tmp/matcreator-cmd-{job_id}" + log_path = f"{marker}.log" + exit_path = f"{marker}.exit" + payload = base64.b64encode(command.encode("utf-8")).decode("ascii") + launch = ( + f"rm -f {exit_path}; " + f"nohup sh -c 'echo {payload} | base64 -d | sh; echo $? > {exit_path}' " + f"> {log_path} 2>&1 < /dev/null & echo LAUNCHED" + ) + launch_result = adapter.run_command(job["external_id"], launch, user=user) + handle = {"log_path": log_path, "exit_path": exit_path, "started_at": time.time()} + self.store.merge_observation( + job_id, + snapshot={"provider_status": "reachable", "background_command": handle}, + error=None, + ) + return {"job_id": job_id, "launch": launch_result, "handle": handle} + + def poll_job_command(self, job_id: str, *, tail_bytes: int = 8000) -> dict[str, Any]: + """Check on the job's most recently started background command. + + Reads only the durable marker-file paths from the job's snapshot, so + this works identically whether it's the same process that started + the command or a freshly re-attached one after a restart. Returns + ``{"running": True, ...}`` while the exit marker hasn't appeared yet, + or ``{"running": False, "exit_code": ..., "output_tail": ...}`` once + it has — ``output_tail`` is the last ``tail_bytes`` of combined + stdout/stderr; use ``download_job_file`` on ``log_path`` for the + full output of a long-running command. + """ + job = self._get_job(job_id) + adapter = self._adapter(job["provider"]) + if RemoteJobCapability.INTERACTIVE_EXEC not in adapter.capabilities: + raise CapabilityError(job["provider"], RemoteJobCapability.INTERACTIVE_EXEC) + handle = (job.get("snapshot") or {}).get("background_command") + if not isinstance(handle, dict) or not handle.get("exit_path"): + raise ValueError(f"Job '{job_id}' has no in-flight background command") + + exit_path = handle["exit_path"] + log_path = handle["log_path"] + check = adapter.run_command( + job["external_id"], + f"if [ -f {exit_path} ]; then echo DONE:$(cat {exit_path}); else echo RUNNING; fi", + user="root", + ) + stdout = str(check.get("stdout", "")).strip() + if not stdout.startswith("DONE:"): + self.store.merge_observation( + job_id, snapshot={"provider_status": "reachable"}, error=None + ) + return {"running": True, "log_path": log_path} + + try: + exit_code = int(stdout.split(":", 1)[1].strip()) + except (IndexError, ValueError): + exit_code = None + tail = adapter.run_command( + job["external_id"], f"tail -c {int(tail_bytes)} {log_path} 2>/dev/null || true", user="root" + ) + self.store.merge_observation( + job_id, + snapshot={"provider_status": "reachable", "background_command": None, "last_command_exit_code": exit_code}, + error=None, + ) + return { + "running": False, + "exit_code": exit_code, + "output_tail": tail.get("stdout", ""), + "log_path": log_path, + } + + def upload_job_file(self, job_id: str, source: str | Path, destination: str) -> dict[str, Any]: + """Upload one local input file into a tracked interactive job.""" + job = self._get_job(job_id) if job["status"] not in {"queued", "running", "resuming"}: - raise ValueError(f"E2B job '{job_id}' cannot receive files while {job['status']}") + raise ValueError(f"Job '{job_id}' cannot receive files while {job['status']}") + adapter = self._adapter(job["provider"]) + if RemoteJobCapability.FILE_TRANSFER not in adapter.capabilities: + raise CapabilityError(job["provider"], RemoteJobCapability.FILE_TRANSFER) source_path = Path(source).expanduser().resolve() - self.e2b_adapter.upload_file(job["external_id"], source_path, destination) + adapter.upload_file(job["external_id"], source_path, destination) self.store.merge_observation( job_id, snapshot={"provider_status": "reachable", "last_upload": source_path.name}, @@ -190,17 +391,16 @@ def upload_e2b_file(self, job_id: str, source: str | Path, destination: str) -> ) return {"source": str(source_path), "destination": destination} - def download_e2b_file(self, job_id: str, source: str, destination: str | Path) -> dict[str, Any]: - """Download one sandbox file to a local destination path. - - Streams the file via the E2B filesystem API so large outputs (CHGCAR, - vasprun.xml, PNG) are not truncated by command-output limits. - """ - job = self._get_e2b_job(job_id) + def download_job_file(self, job_id: str, source: str, destination: str | Path) -> dict[str, Any]: + """Download one file from a tracked interactive job to a local path.""" + job = self._get_job(job_id) if job["status"] not in {"queued", "running", "resuming"}: - raise ValueError(f"E2B job '{job_id}' cannot serve files while {job['status']}") + raise ValueError(f"Job '{job_id}' cannot serve files while {job['status']}") + adapter = self._adapter(job["provider"]) + if RemoteJobCapability.FILE_TRANSFER not in adapter.capabilities: + raise CapabilityError(job["provider"], RemoteJobCapability.FILE_TRANSFER) dest_path = Path(destination).expanduser().resolve() - self.e2b_adapter.download_file(job["external_id"], source, dest_path) + adapter.download_file(job["external_id"], source, dest_path) self.store.merge_observation( job_id, snapshot={"provider_status": "reachable", "last_download": Path(source).name}, @@ -208,12 +408,35 @@ def download_e2b_file(self, job_id: str, source: str, destination: str | Path) - ) return {"source": source, "destination": str(dest_path)} - def _get_e2b_job(self, job_id: str) -> dict[str, Any]: - job = self.store.get_job(job_id) - if job is None: - raise KeyError(f"Remote job '{job_id}' was not found") - if job["provider"] != "e2b": - raise ValueError(f"Remote job '{job_id}' is not managed by E2B") - if not job["external_id"]: - raise ValueError(f"Remote job '{job_id}' has no sandbox ID") - return job \ No newline at end of file + def collect_job_outputs(self, job_id: str, destination_dir: str | Path) -> dict[str, Any]: + """Pull a finished batch job's output files into ``destination_dir``. + + Only valid once the job has reached ``succeeded``; transitions the + job through ``collecting`` -> ``collected`` so a repeated collection + request is a durable no-op rather than a duplicate download. + """ + job = self._get_job(job_id) + if job["status"] == "collected": + return job + if job["status"] != "succeeded": + raise ValueError(f"Job '{job_id}' cannot collect outputs while {job['status']}") + adapter = self._adapter(job["provider"]) + if RemoteJobCapability.BATCH_COLLECT not in adapter.capabilities: + raise CapabilityError(job["provider"], RemoteJobCapability.BATCH_COLLECT) + collecting = self.store.transition_job(job_id, "collecting") + dest_path = Path(destination_dir).expanduser().resolve() + try: + artifacts = adapter.collect_outputs(job["external_id"], dest_path) + except Exception as exc: + return self.store.transition_job( + job_id, + "failed", + error=f"{job['provider']} output collection failed: {exc}", + expected_revision=collecting["state_revision"], + ) + return self.store.transition_job( + job_id, + "collected", + artifacts=artifacts, + expected_revision=collecting["state_revision"], + ) diff --git a/src/matcreator/control_plane/remote_jobs.py b/src/matcreator/control_plane/remote_jobs.py index e4a2c621..b736ea50 100644 --- a/src/matcreator/control_plane/remote_jobs.py +++ b/src/matcreator/control_plane/remote_jobs.py @@ -265,6 +265,47 @@ def transition_job( ) return self.get_job(job_id) or {} + def reset_failed_job_for_retry(self, job_id: str) -> dict[str, Any]: + """Return a failed job that never acquired an external ID to ``created``. + + ``failed`` is terminal for the normal transition machinery, but a job + that failed before the provider handed back an external ID has no + provider-side effect to duplicate, so re-running its submission is + safe. This is the one sanctioned exception, recorded as its own + ``retry`` event. Raises ``ValueError`` for any other job state. + """ + now = time.time() + with self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + row = connection.execute("SELECT * FROM remote_jobs WHERE job_id = ?", (job_id,)).fetchone() + if row is None: + raise KeyError(f"Remote job '{job_id}' was not found") + current = self._decode(row) or {} + if current["status"] != "failed" or current["external_id"]: + raise ValueError( + f"Remote job '{job_id}' cannot be reset for retry " + f"(status={current['status']!r}, external_id={current['external_id']!r})" + ) + updated = connection.execute( + """ + UPDATE remote_jobs + SET status = 'created', error = NULL, + state_revision = state_revision + 1, updated_at = ? + WHERE job_id = ? AND state_revision = ? + """, + (now, job_id, current["state_revision"]), + ) + if updated.rowcount != 1: + raise RuntimeError("Remote job revision changed") + self._append_event( + connection, + job_id, + "retry", + {"from": "failed", "to": "created", "previous_error": current["error"]}, + now, + ) + return self.get_job(job_id) or {} + def record_observation( self, job_id: str, diff --git a/src/matcreator/skills/e2b/SKILL.md b/src/matcreator/skills/e2b/SKILL.md index 2023fe80..8bcf36f9 100644 --- a/src/matcreator/skills/e2b/SKILL.md +++ b/src/matcreator/skills/e2b/SKILL.md @@ -18,9 +18,7 @@ control the sandbox even after the agent or browser reconnects. ## Submission -1. Choose `template` explicitly for every `submit_e2b_sandbox` call. When the - template name is unknown, run `lbg sdbx template ls -q` to list available - templates. Install the command with `pip install -U --pre lbg` when needed. +1. Choose `template` explicitly for every `submit_e2b_sandbox` call. Ask for the `template` explicitly if unknown. 2. Call `submit_e2b_sandbox` once for the current step. It is idempotent for the current session, node, and template. 3. Use `upload_e2b_input` for workspace files, then use `run_e2b_command` for diff --git a/src/matcreator/skills/vasp-pymatgen/references/e2b-sandbox-execution.md b/src/matcreator/skills/vasp-pymatgen/references/e2b-sandbox-execution.md index dd65826a..9e8e458b 100644 --- a/src/matcreator/skills/vasp-pymatgen/references/e2b-sandbox-execution.md +++ b/src/matcreator/skills/vasp-pymatgen/references/e2b-sandbox-execution.md @@ -100,7 +100,7 @@ export E2B_VALIDATE_API_KEY=false ``` The adapter already sets `E2B_VALIDATE_API_KEY=false` by default at import -time (`matcreator/control_plane/e2b.py`), so this is handled for Bohrium +time (`matcreator/control_plane/providers/e2b.py`), so this is handled for Bohrium deployments out of the box. An explicit value in the environment takes precedence — set `E2B_VALIDATE_API_KEY=true` only if you are talking to the public e2b.dev API with a standard `e2b_`-prefixed key. diff --git a/tests/test_bohr_job_adapter.py b/tests/test_bohr_job_adapter.py new file mode 100644 index 00000000..3f91dccb --- /dev/null +++ b/tests/test_bohr_job_adapter.py @@ -0,0 +1,159 @@ +from __future__ import annotations + +import json +import subprocess + +import pytest + +from matcreator.control_plane.providers._bohr_cli import BohrCLIError +from matcreator.control_plane.providers.base import RemoteJobCapability +from matcreator.control_plane.providers.bohr_job import BohrJobAdapter + + +class _FakeCompleted: + def __init__(self, stdout: str, returncode: int = 0, stderr: str = "") -> None: + self.stdout = stdout + self.stderr = stderr + self.returncode = returncode + + +def _ok(data) -> str: + return json.dumps({"ok": True, "data": data}) + + +def _err(message: str) -> str: + return json.dumps({"ok": False, "error": {"message": message}}) + + +def test_create_submits_job_and_extracts_bohr_id(monkeypatch) -> None: + captured = {} + + def fake_run(command, **kwargs): + captured["command"] = command + return _FakeCompleted(_ok({"bohrId": 20543207, "id": 23197091})) + + monkeypatch.setattr(subprocess, "run", fake_run) + adapter = BohrJobAdapter() + + external_id = adapter.create( + { + "project_id": 42, + "job_name": "relax-job", + "machine_type": "c8_m32_cpu", + "image_address": "registry.dp.tech/dptech/vasp:5.4.4", + "command": "vasp_std", + "input_directory": "./input", + "max_run_time": 60, + } + ) + + assert external_id == "20543207" + command = captured["command"] + assert command[1:4] == ["job", "submit", "--project_id"] + assert "42" in command + assert "--input_directory" in command and "./input" in command + assert "--max_run_time" in command and "60" in command + assert "-o" in command and "json" in command + assert "--no-interactive" in command + assert "-y" in command + + +def test_create_requires_all_fields() -> None: + adapter = BohrJobAdapter() + + with pytest.raises(ValueError, match="job_name"): + adapter.create({"project_id": 1}) + + +def test_status_maps_phase_to_normalized_status(monkeypatch) -> None: + def fake_run(command, **kwargs): + return _FakeCompleted(_ok({"phase": "completed", "terminal": True})) + + monkeypatch.setattr(subprocess, "run", fake_run) + adapter = BohrJobAdapter() + + status = adapter.status("20543207") + + assert status.normalized_status == "succeeded" + assert status.snapshot == {"phase": "completed", "terminal": True} + assert status.error is None + + +def test_status_maps_failed_phase_and_captures_error_info(monkeypatch) -> None: + def fake_run(command, **kwargs): + return _FakeCompleted(_ok({"phase": "failed", "terminal": True, "errorInfo": "Command not found."})) + + monkeypatch.setattr(subprocess, "run", fake_run) + adapter = BohrJobAdapter() + + status = adapter.status("20543187") + + assert status.normalized_status == "failed" + assert status.error == "Command not found." + + +def test_status_maps_stopped_phase_to_cancelled(monkeypatch) -> None: + def fake_run(command, **kwargs): + return _FakeCompleted(_ok({"phase": "stopped", "terminal": True})) + + monkeypatch.setattr(subprocess, "run", fake_run) + adapter = BohrJobAdapter() + + status = adapter.status("20539094") + + assert status.normalized_status == "cancelled" + + +def test_status_raises_bohr_cli_error_on_failure(monkeypatch) -> None: + def fake_run(command, **kwargs): + return _FakeCompleted(_err("record not found")) + + monkeypatch.setattr(subprocess, "run", fake_run) + adapter = BohrJobAdapter() + + with pytest.raises(BohrCLIError, match="record not found"): + adapter.status("nonexistent") + + +def test_cancel_invokes_terminate_with_no_wait(monkeypatch) -> None: + captured = {} + + def fake_run(command, **kwargs): + captured["command"] = command + return _FakeCompleted(_ok(None)) + + monkeypatch.setattr(subprocess, "run", fake_run) + adapter = BohrJobAdapter() + + adapter.cancel("20543207") + + command = captured["command"] + assert command[1:5] == ["job", "terminate", "--id", "20543207"] + assert "--no-wait" in command + + +def test_collect_outputs_downloads_and_lists_files(monkeypatch, tmp_path) -> None: + def fake_run(command, **kwargs): + # Simulate the download command producing files in the destination. + dest = tmp_path / "out" + dest.mkdir(parents=True, exist_ok=True) + (dest / "OUTCAR").write_text("data") + (dest / "vasprun.xml").write_text("data") + return _FakeCompleted(_ok(None)) + + monkeypatch.setattr(subprocess, "run", fake_run) + adapter = BohrJobAdapter() + + artifacts = adapter.collect_outputs("20543207", tmp_path / "out") + + sources = {artifact["source"] for artifact in artifacts} + destinations = {artifact["destination"] for artifact in artifacts} + assert sources == {"20543207"} + assert destinations == {str(tmp_path / "out" / "OUTCAR"), str(tmp_path / "out" / "vasprun.xml")} + + +def test_capabilities_are_batch_collect_only() -> None: + adapter = BohrJobAdapter() + + assert adapter.capabilities == frozenset({RemoteJobCapability.BATCH_COLLECT}) + assert adapter.provider == "bohr_job" diff --git a/tests/test_bohr_sandbox_adapter.py b/tests/test_bohr_sandbox_adapter.py new file mode 100644 index 00000000..3c541935 --- /dev/null +++ b/tests/test_bohr_sandbox_adapter.py @@ -0,0 +1,240 @@ +from __future__ import annotations + +import json +import subprocess + +import pytest + +from matcreator.control_plane.providers._bohr_cli import BohrCLIError +from matcreator.control_plane.providers.base import RemoteJobCapability +from matcreator.control_plane.providers.bohr_sandbox import BohrSandboxAdapter + + +class _FakeCompleted: + def __init__(self, stdout: str, returncode: int = 0, stderr: str = "") -> None: + self.stdout = stdout + self.stderr = stderr + self.returncode = returncode + + +def _ok(data) -> str: + return json.dumps({"ok": True, "data": data}) + + +def _err(message: str) -> str: + return json.dumps({"ok": False, "error": {"message": message}}) + + +def test_create_builds_command_and_extracts_sandbox_id(monkeypatch) -> None: + captured = {} + + def fake_run(command, **kwargs): + captured["command"] = command + # Real `bohr sandbox create -o json` payload shape (CLI 2.6.15): + # the ID key is "sandboxID", not "sandbox_id". + return _FakeCompleted( + _ok( + { + "sandboxID": "default--sdbxdefault-abc12", + "templateID": "sdbxagent", + "state": "running", + "domain": "bohr-sandbox.bohrium.com", + } + ) + ) + + monkeypatch.setattr(subprocess, "run", fake_run) + adapter = BohrSandboxAdapter() + + sandbox_id = adapter.create( + { + "project_id": 1234, + "template": "sdbxagent", + "timeout": 3600, + "env": {"FOO": "bar"}, + } + ) + + assert sandbox_id == "default--sdbxdefault-abc12" + command = captured["command"] + assert command[1:3] == ["sandbox", "create"] + assert "--template" in command and "sdbxagent" in command + assert "--project-id" in command and "1234" in command + assert "--timeout" in command and "3600" in command + assert "--env" in command and "FOO=bar" in command + + +def test_create_requires_template() -> None: + adapter = BohrSandboxAdapter() + + with pytest.raises(ValueError, match="template"): + adapter.create({"project_id": 1234}) + + +def test_create_accepts_snake_case_sandbox_id_spelling(monkeypatch) -> None: + def fake_run(command, **kwargs): + return _FakeCompleted(_ok({"sandbox_id": "default--sdbxdefault-abc12"})) + + monkeypatch.setattr(subprocess, "run", fake_run) + + assert BohrSandboxAdapter().create( + {"project_id": 1, "template": "sdbxagent"} + ) == "default--sdbxdefault-abc12" + + +def test_create_requires_project_id() -> None: + adapter = BohrSandboxAdapter() + + with pytest.raises(ValueError, match="project_id"): + adapter.create({}) + + +def test_create_raises_when_no_sandbox_id_returned(monkeypatch) -> None: + def fake_run(command, **kwargs): + return _FakeCompleted(_ok({"state": "running", "templateID": "sdbxagent"})) + + monkeypatch.setattr(subprocess, "run", fake_run) + adapter = BohrSandboxAdapter() + + with pytest.raises(BohrCLIError, match=r"did not return a sandbox ID.*state.*templateID"): + adapter.create({"project_id": 1, "template": "sdbxagent"}) + + +def test_status_reports_liveness_without_normalized_status_by_default(monkeypatch) -> None: + def fake_run(command, **kwargs): + return _FakeCompleted(_ok({"status": "running"})) + + monkeypatch.setattr(subprocess, "run", fake_run) + adapter = BohrSandboxAdapter() + + status = adapter.status("sbx-1") + + assert status.normalized_status is None + assert status.snapshot["provider_status"] == "reachable" + + +def test_status_maps_unreachable_states_to_lost(monkeypatch) -> None: + def fake_run(command, **kwargs): + return _FakeCompleted(_ok({"status": "terminated"})) + + monkeypatch.setattr(subprocess, "run", fake_run) + adapter = BohrSandboxAdapter() + + status = adapter.status("sbx-1") + + assert status.normalized_status == "lost" + assert status.snapshot["provider_status"] == "unreachable" + + +def test_cancel_invokes_delete_with_force(monkeypatch) -> None: + captured = {} + + def fake_run(command, **kwargs): + captured["command"] = command + return _FakeCompleted(_ok(None)) + + monkeypatch.setattr(subprocess, "run", fake_run) + adapter = BohrSandboxAdapter() + + adapter.cancel("sbx-1") + + command = captured["command"] + assert command[1:4] == ["sandbox", "delete", "sbx-1"] + assert "--force" in command + + +def test_run_command_parses_stdout_stderr_exit_code(monkeypatch) -> None: + captured = {} + + def fake_run(command, **kwargs): + captured["command"] = command + captured["kwargs"] = kwargs + return _FakeCompleted(_ok({"stdout": "hello\n", "stderr": "", "exit_code": 0})) + + monkeypatch.setattr(subprocess, "run", fake_run) + adapter = BohrSandboxAdapter() + + result = adapter.run_command("sbx-1", "echo hello") + + assert result == {"stdout": "hello\n", "stderr": "", "exit_code": 0} + command = captured["command"] + assert command[1:5] == ["sandbox", "exec", "sbx-1", "--command"] + assert "echo hello" in command + + +def test_run_command_disables_both_cli_and_subprocess_timeouts(monkeypatch) -> None: + """`bohr sandbox exec` caps a command at 90s by default; a long-running + remote computation must not be silently truncated, matching the E2B + adapter's unbounded `timeout=0` command semantics.""" + captured = {} + + def fake_run(command, **kwargs): + captured["command"] = command + captured["kwargs"] = kwargs + return _FakeCompleted(_ok({"stdout": "", "stderr": "", "exit_code": 0})) + + monkeypatch.setattr(subprocess, "run", fake_run) + adapter = BohrSandboxAdapter() + + adapter.run_command("sbx-1", "sleep 300") + + command = captured["command"] + timeout_index = command.index("--timeout") + assert command[timeout_index + 1] == "0" + assert captured["kwargs"]["timeout"] is None + + +def test_upload_file_rejects_missing_source(tmp_path) -> None: + adapter = BohrSandboxAdapter() + + with pytest.raises(FileNotFoundError): + adapter.upload_file("sbx-1", tmp_path / "missing.txt", "/home/user/missing.txt") + + +def test_upload_file_invokes_files_write(monkeypatch, tmp_path) -> None: + source = tmp_path / "input.txt" + source.write_text("hello") + captured = {} + + def fake_run(command, **kwargs): + captured["command"] = command + return _FakeCompleted(_ok(None)) + + monkeypatch.setattr(subprocess, "run", fake_run) + adapter = BohrSandboxAdapter() + + adapter.upload_file("sbx-1", source, "/home/user/input.txt") + + command = captured["command"] + assert command[1:5] == ["sandbox", "files", "write", "sbx-1"] + assert "/home/user/input.txt" in command + assert "--source" in command and str(source) in command + + +def test_download_file_invokes_files_read_and_creates_parent_dir(monkeypatch, tmp_path) -> None: + captured = {} + + def fake_run(command, **kwargs): + captured["command"] = command + return _FakeCompleted(_ok(None)) + + monkeypatch.setattr(subprocess, "run", fake_run) + adapter = BohrSandboxAdapter() + dest = tmp_path / "outputs" / "CHGCAR" + + result = adapter.download_file("sbx-1", "/home/user/CHGCAR", dest) + + assert result == dest.resolve() + assert dest.parent.is_dir() + command = captured["command"] + assert command[1:5] == ["sandbox", "files", "read", "sbx-1"] + assert "--destination" in command and str(dest.resolve()) in command + + +def test_capabilities_have_no_pause() -> None: + adapter = BohrSandboxAdapter() + + assert adapter.capabilities == frozenset( + {RemoteJobCapability.INTERACTIVE_EXEC, RemoteJobCapability.FILE_TRANSFER} + ) + assert adapter.provider == "bohr_sandbox" diff --git a/tests/test_e2b_adapter.py b/tests/test_e2b_adapter.py index ed363c8e..b3a43cbe 100644 --- a/tests/test_e2b_adapter.py +++ b/tests/test_e2b_adapter.py @@ -5,7 +5,7 @@ import pytest -from matcreator.control_plane.e2b import E2BConfigurationError, E2BSandboxAdapter, E2BSandboxSpec +from matcreator.control_plane.providers.e2b import E2BConfigurationError, E2BSandboxAdapter, E2BSandboxSpec class _FakeResult: @@ -52,6 +52,7 @@ class _FakeSandbox: sandbox_id = "sandbox-123" created_with: dict = {} connected_to: list[str] = [] + connect_opts: list[dict] = [] paused = False killed = False files = _FakeFiles() @@ -62,14 +63,17 @@ def create(cls, **kwargs): return cls() @classmethod - def connect(cls, sandbox_id): + def connect(cls, sandbox_id, **opts): cls.connected_to.append(sandbox_id) + cls.connect_opts.append(opts) return cls() class commands: + last_command: str | None = None + @staticmethod def run(command, user, **kwargs): - assert command == "echo hello" + _FakeSandbox.commands.last_command = command assert user == "root" return _FakeResult() @@ -84,22 +88,27 @@ def kill(self): def fake_e2b_module(monkeypatch): _FakeSandbox.created_with = {} _FakeSandbox.connected_to = [] + _FakeSandbox.connect_opts = [] _FakeSandbox.paused = False _FakeSandbox.killed = False _FakeSandbox.files = _FakeFiles() monkeypatch.setitem(sys.modules, "e2b_code_interpreter", types.SimpleNamespace(Sandbox=_FakeSandbox)) + # Reconnects require the endpoint configuration in the environment. + monkeypatch.setenv("E2B_API_KEY", "secret") + monkeypatch.setenv("E2B_API_URL", "https://e2b.example") + monkeypatch.setenv("BOHRIUM_PROJECT_ID", "project-42") def test_adapter_creates_sandbox_with_project_header() -> None: adapter = E2BSandboxAdapter() sandbox_id = adapter.create( - E2BSandboxSpec( - template="doc-compiler", - api_key="secret", - api_url="https://e2b.example", - project_id="project-42", - lifecycle={"on_timeout": "pause"}, - ) + { + "template": "doc-compiler", + "api_key": "secret", + "api_url": "https://e2b.example", + "project_id": "project-42", + "lifecycle": {"on_timeout": "pause"}, + } ) assert sandbox_id == "sandbox-123" @@ -115,12 +124,70 @@ def test_adapter_connects_for_command_and_controls() -> None: "stderr": "", "exit_code": 0, } + assert _FakeSandbox.commands.last_command == "echo hello" adapter.pause("sandbox-123") adapter.terminate("sandbox-123") assert _FakeSandbox.connected_to == ["sandbox-123", "sandbox-123", "sandbox-123"] assert _FakeSandbox.paused is True assert _FakeSandbox.killed is True + # Reconnects must carry the configured endpoint, never the SDK default. + assert _FakeSandbox.connect_opts[0] == { + "api_key": "secret", + "api_url": "https://e2b.example", + "headers": {"X-Project-Id": "project-42"}, + } + + +def test_adapter_connect_fails_loudly_without_endpoint_configuration(monkeypatch) -> None: + adapter = E2BSandboxAdapter() + monkeypatch.delenv("E2B_API_KEY", raising=False) + + with pytest.raises(E2BConfigurationError, match="E2B_API_KEY and E2B_API_URL"): + adapter.run_command("sandbox-123", "echo hello") + assert _FakeSandbox.connected_to == [] + + +def test_adapter_cancel_aliases_terminate() -> None: + adapter = E2BSandboxAdapter() + + adapter.cancel("sandbox-123") + + assert _FakeSandbox.killed is True + + +def test_adapter_create_error_includes_request_context(monkeypatch) -> None: + def _boom(**kwargs): + raise RuntimeError("404: Resource not found") + + monkeypatch.setattr(_FakeSandbox, "create", _boom) + adapter = E2BSandboxAdapter() + + with pytest.raises(RuntimeError) as excinfo: + adapter.create( + { + "template": "doc-compiler", + "api_key": "secret", + "api_url": "https://open.bohrium.com/wrong/path", + "project_id": "project-42", + } + ) + + message = str(excinfo.value) + assert "404: Resource not found" in message + assert "https://open.bohrium.com/wrong/path" in message + assert "doc-compiler" in message + assert "project-42" in message + assert "secret" not in message + + +def test_adapter_status_reports_liveness_without_a_normalized_status() -> None: + adapter = E2BSandboxAdapter() + + status = adapter.status("sandbox-123") + + assert status.normalized_status is None + assert status.snapshot["provider_status"] == "reachable" def test_adapter_download_file_streams_to_local_destination(tmp_path) -> None: diff --git a/tests/test_e2b_tools.py b/tests/test_e2b_tools.py deleted file mode 100644 index 75839cea..00000000 --- a/tests/test_e2b_tools.py +++ /dev/null @@ -1,184 +0,0 @@ -from __future__ import annotations - -from types import SimpleNamespace - -from matcreator.agents.execution_agent import e2b_tools -from matcreator.control_plane.remote_job_service import E2BConnectionConfig - - -class _FakeService: - def __init__(self) -> None: - self.submissions: list[dict] = [] - self.store = self - - def submit_e2b(self, **kwargs): - self.submissions.append(kwargs) - return { - "job_id": "job-123", - "status": "running", - "external_id": "sandbox-123", - } - - def get_job(self, job_id: str): - if job_id != "job-123": - return None - return { - "job_id": job_id, - "owner_id": "alice", - "session_id": "session-1", - "status": "running", - "external_id": "sandbox-123", - "snapshot": {}, - "error": None, - "updated_at": 1, - } - - def list_events(self, job_id: str): - return [{"event_type": "user_control", "payload": {"action": "terminate", "source": "ui"}}] - - def pause_e2b(self, job_id: str): - return {"job_id": job_id, "status": "paused", "external_id": "sandbox-123"} - - def terminate_e2b(self, job_id: str): - return {"job_id": job_id, "status": "terminated", "external_id": "sandbox-123"} - - def run_e2b_command(self, job_id: str, command: str, *, user: str): - return {"stdout": f"ran {command}", "stderr": "", "exit_code": 0} - - def upload_e2b_file(self, job_id: str, source, destination: str): - return {"source": str(source), "destination": destination} - - def download_e2b_file(self, job_id: str, source: str, destination: str): - return {"source": source, "destination": str(destination)} - - -def _context(): - return SimpleNamespace( - state={ - "session_id": "session-1", - "_graph_exec_node_id": "execution_0__node_relax", - "step_number": 2, - }, - _invocation_context=SimpleNamespace(user_id="alice"), - ) - - -def test_submit_e2b_tool_uses_current_session_and_node(monkeypatch) -> None: - service = _FakeService() - monkeypatch.setattr(e2b_tools, "_service", lambda: service) - monkeypatch.setenv("E2B_API_KEY", "secret") - monkeypatch.setenv("E2B_API_URL", "https://e2b.example") - monkeypatch.setenv("BOHRIUM_PROJECT_ID", "project-42") - - result = e2b_tools.submit_e2b_sandbox(_context(), timeout=120, template="doc-compiler") - - assert result == { - "status": "running", - "job_id": "job-123", - "sandbox_id": "sandbox-123", - "message": "Tracked E2B sandbox is ready. Use its job_id for status or controls.", - } - submission = service.submissions[0] - assert submission["owner_id"] == "alice" - assert submission["session_id"] == "session-1" - assert submission["node_id"] == "relax" - assert submission["step_number"] == 2 - assert submission["connection"] == E2BConnectionConfig( - api_key="secret", - api_url="https://e2b.example", - project_id="project-42", - template="doc-compiler", - ) - - -def test_submit_e2b_tool_requires_explicit_template(monkeypatch) -> None: - service = _FakeService() - monkeypatch.setattr(e2b_tools, "_service", lambda: service) - - result = e2b_tools.submit_e2b_sandbox(_context()) - - assert result["status"] == "error" - assert "template is required" in result["message"] - assert service.submissions == [] - - -def test_e2b_connection_uses_configured_environment_names(monkeypatch) -> None: - monkeypatch.setenv("E2B_API_KEY", "access-key") - monkeypatch.setenv("E2B_API_URL", "https://e2b.example") - monkeypatch.setenv("BOHRIUM_PROJECT_ID", "project-7") - - connection = e2b_tools._connection() - - assert connection == E2BConnectionConfig( - api_key="access-key", - api_url="https://e2b.example", - project_id="project-7", - template="", - ) - - -def test_e2b_tools_reject_jobs_from_another_session(monkeypatch) -> None: - service = _FakeService() - monkeypatch.setattr(e2b_tools, "_service", lambda: service) - context = _context() - context._invocation_context.user_id = "bob" - - assert e2b_tools.get_e2b_job_status("job-123", context) == { - "status": "error", - "message": "E2B job was not found in this session.", - } - - -def test_e2b_status_exposes_user_sandbox_control(monkeypatch) -> None: - service = _FakeService() - monkeypatch.setattr(e2b_tools, "_service", lambda: service) - - status = e2b_tools.get_e2b_job_status("job-123", _context()) - - assert status["user_control"] == {"action": "terminate", "source": "ui"} - - -def test_e2b_command_and_workspace_upload_are_scoped_to_owned_job(tmp_path, monkeypatch) -> None: - service = _FakeService() - monkeypatch.setattr(e2b_tools, "_service", lambda: service) - context = _context() - context.state["workspace_dir"] = str(tmp_path) - source = tmp_path / "input.txt" - source.write_text("input", encoding="utf-8") - - assert e2b_tools.run_e2b_command("job-123", "echo hello", context) == { - "stdout": "ran echo hello", "stderr": "", "exit_code": 0 - } - assert e2b_tools.upload_e2b_input("job-123", "input.txt", "/home/user/input.txt", context) == { - "source": str(source), "destination": "/home/user/input.txt" - } - assert e2b_tools.upload_e2b_input("job-123", "/tmp/outside.txt", "/tmp/outside.txt", context)["status"] == "error" - - -def test_download_e2b_output_is_scoped_to_workspace(tmp_path, monkeypatch) -> None: - service = _FakeService() - monkeypatch.setattr(e2b_tools, "_service", lambda: service) - context = _context() - context.state["workspace_dir"] = str(tmp_path) - destination = tmp_path / "outputs" / "CHGCAR" - - result = e2b_tools.download_e2b_output( - "job-123", "/home/user/CHGCAR", str(destination), context - ) - assert result == {"source": "/home/user/CHGCAR", "destination": str(destination.resolve())} - - assert e2b_tools.download_e2b_output( - "job-123", "/home/user/CHGCAR", "/tmp/outside.txt", context - )["status"] == "error" - - -def test_download_e2b_output_rejects_missing_workspace_dir(monkeypatch) -> None: - service = _FakeService() - monkeypatch.setattr(e2b_tools, "_service", lambda: service) - context = _context() # no workspace_dir set - - result = e2b_tools.download_e2b_output( - "job-123", "/home/user/CHGCAR", "outputs/CHGCAR", context - ) - assert result["status"] == "error" - assert "workspace_dir" in result["message"] \ No newline at end of file diff --git a/tests/test_execution_grouping.py b/tests/test_execution_grouping.py new file mode 100644 index 00000000..cc8943b4 --- /dev/null +++ b/tests/test_execution_grouping.py @@ -0,0 +1,129 @@ +import unittest + +from pydantic import ValidationError + +from agents.MatCreator.agents.execution_agent.agent import ( + build_execution_groups, + build_execution_waves, +) +from agents.MatCreator.agents.execution_agent.step_executor import StepExecutorInput + + +class _DummyToolContext: + def __init__(self, state: dict): + self.state = state + + +class TestExecutionGrouping(unittest.TestCase): + def test_groups_consecutive_same_skill(self) -> None: + ctx = _DummyToolContext( + { + "current_step_index": 0, + "plan": { + "steps": [ + {"step_number": 1, "skill": "vasp", "action": "Prepare input files."}, + {"step_number": 2, "skill": "vasp", "action": "Run relaxation."}, + {"step_number": 3, "skill": "plot", "action": "Plot total energy."}, + ] + }, + } + ) + + result = build_execution_groups(ctx) + self.assertEqual(result["status"], "ok") + self.assertEqual(len(result["groups"]), 2) + self.assertEqual(result["groups"][0]["step_numbers"], [1, 2]) + self.assertEqual(result["groups"][0]["skill_name"], "vasp") + self.assertEqual(result["groups"][1]["step_numbers"], [3]) + self.assertEqual(result["groups"][1]["skill_name"], "plot") + + def test_dependency_marker_starts_new_group(self) -> None: + ctx = _DummyToolContext( + { + "current_step_index": 0, + "plan": { + "steps": [ + {"step_number": 1, "skill": "vasp", "action": "Run static calculation."}, + { + "step_number": 2, + "skill": "vasp", + "action": "Using previous step results, extract DOS.", + }, + ] + }, + } + ) + + result = build_execution_groups(ctx) + self.assertEqual(result["status"], "ok") + self.assertEqual(len(result["groups"]), 2) + self.assertEqual(result["groups"][0]["step_numbers"], [1]) + self.assertEqual(result["groups"][1]["step_numbers"], [2]) + + def test_build_execution_waves_parallelizes_distinct_skills(self) -> None: + groups = [ + { + "group_id": "group_1_2", + "skill_name": "vasp", + "step_numbers": [1, 2], + "actions": ["Prepare inputs.", "Run relax."], + }, + { + "group_id": "group_3_3", + "skill_name": "plot", + "step_numbers": [3], + "actions": ["Plot band structure."], + }, + ] + + result = build_execution_waves(groups) + self.assertEqual(result["status"], "ok") + self.assertEqual(len(result["waves"]), 1) + self.assertEqual(len(result["waves"][0]), 2) + + def test_build_execution_waves_serializes_dependency_marked_group(self) -> None: + groups = [ + { + "group_id": "group_1_1", + "skill_name": "vasp", + "step_numbers": [1], + "actions": ["Run SCF."], + }, + { + "group_id": "group_2_2", + "skill_name": "plot", + "step_numbers": [2], + "actions": ["Using previous step results, plot DOS."], + }, + ] + + result = build_execution_waves(groups) + self.assertEqual(result["status"], "ok") + self.assertEqual(len(result["waves"]), 2) + self.assertEqual(result["waves"][0][0]["group_id"], "group_1_1") + self.assertEqual(result["waves"][1][0]["group_id"], "group_2_2") + + +class TestStepExecutorInputNormalization(unittest.TestCase): + def test_legacy_fields_are_normalized(self) -> None: + payload = StepExecutorInput( + step_number=2, + action="Run calculation.", + skill_name="vasp", + workspace_dir="/tmp/work", + ) + self.assertEqual(payload.step_numbers, [2]) + self.assertEqual(payload.actions, ["Run calculation."]) + + def test_mismatched_lengths_fail_validation(self) -> None: + with self.assertRaises(ValidationError): + StepExecutorInput( + step_numbers=[1, 2], + actions=["one"], + skill_name="vasp", + workspace_dir="/tmp/work", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_execution_recovery.py b/tests/test_execution_recovery.py index a41072db..c9a2a862 100644 --- a/tests/test_execution_recovery.py +++ b/tests/test_execution_recovery.py @@ -449,6 +449,7 @@ def test_reconcile_waits_for_active_remote_job_instead_of_resubmitting(tmp_path, "provider": "e2b", "external_id": "sandbox-123", "status": "running", + "has_background_command": False, } diff --git a/tests/test_matcreator_phase_and_skills.py b/tests/test_matcreator_phase_and_skills.py new file mode 100644 index 00000000..285990a2 --- /dev/null +++ b/tests/test_matcreator_phase_and_skills.py @@ -0,0 +1,48 @@ +import unittest + +from agents.MatCreator.agent import before_agent_callback_root +from agents.MatCreator.planning_agent import planning_agent +from agents.MatCreator.prompts.workflow import get_all_workflow_types, search_skills + + +class _DummySession: + def __init__(self): + self.id = "test-session" + self.user_id = "test-user" + self.app_name = "test-app" + self.state = {} + + +class _DummyInvocationContext: + def __init__(self): + self.session = _DummySession() + + +class _DummyCallbackContext: + def __init__(self): + self._invocation_context = _DummyInvocationContext() + + +class TestMatCreatorPhaseAndSkills(unittest.TestCase): + def test_before_agent_callback_sets_default_phase(self) -> None: + callback_context = _DummyCallbackContext() + before_agent_callback_root(callback_context) + self.assertEqual(callback_context._invocation_context.session.state["phase"], "thinking") + + def test_skill_registry_loads_expected_workflows(self) -> None: + workflow_types = get_all_workflow_types() + self.assertIn("default", workflow_types) + self.assertIn("pfd", workflow_types) + + def test_skill_search_returns_matching_workflow(self) -> None: + results = search_skills("fine-tune distillation active learning", workflow_type="pfd", top_k=2) + self.assertGreaterEqual(len(results), 1) + self.assertEqual(results[0].workflow_type, "pfd") + + def test_planning_agent_has_toolized_subagents(self) -> None: + tools = getattr(planning_agent, "tools", []) or [] + self.assertGreaterEqual(len(tools), 3) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_remote_job_monitor.py b/tests/test_remote_job_monitor.py index ced335e1..1ba0e0e7 100644 --- a/tests/test_remote_job_monitor.py +++ b/tests/test_remote_job_monitor.py @@ -2,57 +2,71 @@ import asyncio +from matcreator.control_plane.providers.base import RemoteJobAdapter, RemoteJobCapability, RemoteJobStatus from matcreator.control_plane.remote_job_monitor import RemoteJobMonitor -from matcreator.control_plane.remote_job_service import E2BConnectionConfig, RemoteJobService +from matcreator.control_plane.remote_job_service import RemoteJobService from matcreator.control_plane.remote_jobs import RemoteJobStore -class _FakeE2BAdapter: +class _FakeAdapter(RemoteJobAdapter): + provider = "e2b" + capabilities = frozenset({RemoteJobCapability.PAUSE}) + poll_interval_seconds = 1.0 + def __init__(self, *, reachable: bool = True) -> None: self.reachable = reachable self.probes: list[str] = [] - def create(self, _spec): + def create(self, spec: dict) -> str: return "sandbox-123" - def probe(self, sandbox_id: str): - self.probes.append(sandbox_id) + def status(self, external_id: str) -> RemoteJobStatus: + self.probes.append(external_id) if not self.reachable: raise RuntimeError("sandbox unavailable") - return {"provider_status": "reachable", "sandbox_id": sandbox_id} + return RemoteJobStatus(normalized_status=None, snapshot={"provider_status": "reachable", "sandbox_id": external_id}) + + def cancel(self, external_id: str) -> None: + pass + + def pause(self, external_id: str) -> None: + pass -def _create_running_job(tmp_path, adapter: _FakeE2BAdapter): +def _create_running_job(tmp_path, adapter: _FakeAdapter): store = RemoteJobStore(tmp_path / "remote-jobs.db") - service = RemoteJobService(store, e2b_adapter=adapter) - job = service.submit_e2b( + service = RemoteJobService(store, adapter_overrides={"e2b": adapter}) + job = service.submit_job( owner_id="alice", session_id="session-1", + provider="e2b", idempotency_key="session-1:node-1:1", - connection=E2BConnectionConfig( - api_key="secret", - api_url="https://e2b.example", - project_id="project-42", - template="doc-compiler", - ), + spec={ + "template": "doc-compiler", + "api_key": "secret", + "api_url": "https://e2b.example", + "project_id": "project-42", + }, ) return store, service, job def test_monitor_reconciles_running_job_after_restart(tmp_path) -> None: - adapter = _FakeE2BAdapter() + adapter = _FakeAdapter() store, service, job = _create_running_job(tmp_path, adapter) monitor = RemoteJobMonitor(store, service, interval_seconds=1) updates = asyncio.run(monitor.reconcile_once()) assert [item["job_id"] for item in updates] == [job["job_id"]] - assert adapter.probes == ["sandbox-123"] + # One status() call happens inside submit_job itself (initial probe), one + # more from the explicit reconcile_once() call above. + assert adapter.probes == ["sandbox-123", "sandbox-123"] assert store.get_job(job["job_id"])["snapshot"]["provider_status"] == "reachable" def test_monitor_backs_off_unreachable_job_and_skips_paused_jobs(tmp_path) -> None: - adapter = _FakeE2BAdapter(reachable=False) + adapter = _FakeAdapter(reachable=False) store, service, job = _create_running_job(tmp_path, adapter) monitor = RemoteJobMonitor(store, service, interval_seconds=1, max_backoff_seconds=4) @@ -61,8 +75,53 @@ def test_monitor_backs_off_unreachable_job_and_skips_paused_jobs(tmp_path) -> No assert first[0]["snapshot"]["provider_status"] == "unreachable" assert second == [] - assert adapter.probes == ["sandbox-123"] paused = store.transition_job(job["job_id"], "pause_requested") store.transition_job(job["job_id"], "paused", expected_revision=paused["state_revision"]) monitor._next_due.clear() - assert asyncio.run(monitor.reconcile_once()) == [] \ No newline at end of file + assert asyncio.run(monitor.reconcile_once()) == [] + + +def test_monitor_reconciles_jobs_across_multiple_providers(tmp_path) -> None: + """A batch-style provider with a longer poll interval is reconciled the + same way as an interactive one — the monitor never branches on provider + name, only on each adapter's declared poll_interval_seconds.""" + + class _BatchAdapter(RemoteJobAdapter): + provider = "bohr_job" + capabilities = frozenset({RemoteJobCapability.BATCH_COLLECT}) + poll_interval_seconds = 60.0 + + def __init__(self) -> None: + self.probes: list[str] = [] + + def create(self, spec: dict) -> str: + return "bohr-1" + + def status(self, external_id: str) -> RemoteJobStatus: + self.probes.append(external_id) + return RemoteJobStatus(normalized_status=None, snapshot={"phase": "running"}) + + def cancel(self, external_id: str) -> None: + pass + + e2b_adapter = _FakeAdapter() + batch_adapter = _BatchAdapter() + store = RemoteJobStore(tmp_path / "remote-jobs.db") + service = RemoteJobService(store, adapter_overrides={"e2b": e2b_adapter, "bohr_job": batch_adapter}) + e2b_job = service.submit_job( + owner_id="alice", session_id="session-1", provider="e2b", + idempotency_key="session-1:node-1:1", + spec={"template": "t", "api_key": "k", "api_url": "u", "project_id": "p"}, + ) + batch_job = service.submit_job( + owner_id="alice", session_id="session-1", provider="bohr_job", + idempotency_key="session-1:node-2:1", + spec={"project_id": 1, "job_name": "n", "machine_type": "c2", "image_address": "img", "command": "cmd"}, + ) + + monitor = RemoteJobMonitor(store, service, interval_seconds=1) + updates = asyncio.run(monitor.reconcile_once()) + + reconciled_ids = {item["job_id"] for item in updates} + assert reconciled_ids == {e2b_job["job_id"], batch_job["job_id"]} + assert monitor._next_due[batch_job["job_id"]] > monitor._next_due[e2b_job["job_id"]] diff --git a/tests/test_remote_job_provider_registry.py b/tests/test_remote_job_provider_registry.py new file mode 100644 index 00000000..40fa5a58 --- /dev/null +++ b/tests/test_remote_job_provider_registry.py @@ -0,0 +1,117 @@ +from __future__ import annotations + +import subprocess + +import pytest + +from matcreator.control_plane.providers import registry +from matcreator.control_plane.providers._bohr_cli import BohrCLIError, run_bohr_json +from matcreator.control_plane.providers.base import RemoteJobAdapter, RemoteJobStatus + + +class _FakeCompleted: + def __init__(self, stdout: str = "", stderr: str = "", returncode: int = 0) -> None: + self.stdout = stdout + self.stderr = stderr + self.returncode = returncode + + +def test_run_bohr_json_raises_on_missing_binary(monkeypatch) -> None: + def fake_run(command, **kwargs): + raise FileNotFoundError() + + monkeypatch.setattr(subprocess, "run", fake_run) + + with pytest.raises(BohrCLIError, match="not installed"): + run_bohr_json(["job", "list"]) + + +def test_run_bohr_json_raises_on_timeout(monkeypatch) -> None: + def fake_run(command, **kwargs): + raise subprocess.TimeoutExpired(cmd=command, timeout=1) + + monkeypatch.setattr(subprocess, "run", fake_run) + + with pytest.raises(BohrCLIError, match="timed out"): + run_bohr_json(["job", "list"], timeout=1) + + +def test_run_bohr_json_raises_on_empty_output(monkeypatch) -> None: + def fake_run(command, **kwargs): + return _FakeCompleted(stdout="", stderr="permission denied", returncode=1) + + monkeypatch.setattr(subprocess, "run", fake_run) + + with pytest.raises(BohrCLIError, match="permission denied"): + run_bohr_json(["job", "list"]) + + +def test_run_bohr_json_raises_on_non_json_output(monkeypatch) -> None: + def fake_run(command, **kwargs): + return _FakeCompleted(stdout="not json") + + monkeypatch.setattr(subprocess, "run", fake_run) + + with pytest.raises(BohrCLIError, match="non-JSON"): + run_bohr_json(["job", "list"]) + + +def test_run_bohr_json_returns_data_on_success(monkeypatch) -> None: + def fake_run(command, **kwargs): + return _FakeCompleted(stdout='{"ok": true, "data": {"a": 1}}') + + monkeypatch.setattr(subprocess, "run", fake_run) + + assert run_bohr_json(["job", "list"]) == {"a": 1} + + +class _DummyAdapter(RemoteJobAdapter): + provider = "dummy" + + def create(self, spec): + return "id-1" + + def status(self, external_id): + return RemoteJobStatus(normalized_status=None) + + def cancel(self, external_id): + pass + + +def test_registry_lazy_construction_and_provider_mismatch_detection(): + registry.reset_registry() + try: + calls = [] + + def factory(): + calls.append(1) + return _DummyAdapter() + + registry.register_adapter("dummy", factory) + assert calls == [] # not constructed until first get_adapter call + + adapter = registry.get_adapter("dummy") + assert isinstance(adapter, _DummyAdapter) + assert calls == [1] + + # Cached: second call does not re-invoke the factory. + registry.get_adapter("dummy") + assert calls == [1] + + with pytest.raises(KeyError): + registry.get_adapter("nonexistent") + + class _MismatchedAdapter(_DummyAdapter): + provider = "other-name" + + registry.register_adapter("mismatched", _MismatchedAdapter) + with pytest.raises(ValueError, match="reports provider"): + registry.get_adapter("mismatched") + finally: + registry.reset_registry() + # Re-register built-ins so later tests in the same process still see them. + import importlib + + import matcreator.control_plane.providers as providers_pkg + + importlib.reload(providers_pkg) diff --git a/tests/test_remote_job_service.py b/tests/test_remote_job_service.py index febc81e0..cd740882 100644 --- a/tests/test_remote_job_service.py +++ b/tests/test_remote_job_service.py @@ -1,65 +1,124 @@ from __future__ import annotations -from matcreator.control_plane.remote_job_service import E2BConnectionConfig, RemoteJobService +import base64 +import re + +from matcreator.control_plane.providers.base import RemoteJobAdapter, RemoteJobCapability, RemoteJobStatus +from matcreator.control_plane.remote_job_service import RemoteJobService from matcreator.control_plane.remote_jobs import RemoteJobStore -class _FakeE2BAdapter: +class _FakeAdapter(RemoteJobAdapter): + """Fake adapter conforming to the provider protocol, used via adapter_overrides. + + Declares every optional capability by default so one fake can cover the + submit/pause/terminate/command/upload/download surface; a test can shrink + ``capabilities`` to exercise CapabilityError handling. + """ + + provider = "e2b" + capabilities = frozenset( + { + RemoteJobCapability.PAUSE, + RemoteJobCapability.INTERACTIVE_EXEC, + RemoteJobCapability.FILE_TRANSFER, + } + ) + def __init__(self) -> None: - self.created_specs = [] - self.paused = [] - self.terminated = [] + self.created_specs: list[dict] = [] + self.paused: list[str] = [] + self.cancelled: list[str] = [] self.on_run = None + self.files: dict[str, bytes] = {} + self.launched_commands: list[str] = [] - def create(self, spec): + def create(self, spec: dict) -> str: self.created_specs.append(spec) return "sandbox-123" - def pause(self, sandbox_id: str) -> None: - self.paused.append(sandbox_id) + def status(self, external_id: str) -> RemoteJobStatus: + return RemoteJobStatus(normalized_status=None, snapshot={"provider_status": "reachable"}) + + def cancel(self, external_id: str) -> None: + self.cancelled.append(external_id) - def terminate(self, sandbox_id: str) -> None: - self.terminated.append(sandbox_id) + def pause(self, external_id: str) -> None: + self.paused.append(external_id) - def run_command(self, sandbox_id: str, command: str, *, user: str) -> dict: + def run_command(self, external_id: str, command: str, *, user: str = "root") -> dict: if self.on_run: self.on_run() + # Minimal shell simulation covering the three wrapper patterns + # RemoteJobService.start_job_command/poll_job_command construct, so + # tests can exercise them without a real shell. + launch = re.search( + r"rm -f (\S+); nohup sh -c 'echo (\S+) \| base64 -d \| sh; echo \$\? > \S+' > (\S+) 2>&1", + command, + ) + if launch: + exit_path, payload, log_path = launch.groups() + self.launched_commands.append(base64.b64decode(payload).decode("utf-8")) + self.files.pop(exit_path, None) + self.files.setdefault(log_path, b"") + return {"stdout": "LAUNCHED\n", "stderr": "", "exit_code": 0} + check = re.match(r"if \[ -f (\S+) \]; then echo DONE:\$\(cat \S+\); else echo RUNNING; fi", command) + if check: + exit_path = check.group(1) + if exit_path in self.files: + code = self.files[exit_path].decode("utf-8").strip() + return {"stdout": f"DONE:{code}\n", "stderr": "", "exit_code": 0} + return {"stdout": "RUNNING\n", "stderr": "", "exit_code": 0} + tail = re.match(r"tail -c (\d+) (\S+)", command) + if tail: + n, log_path = tail.groups() + data = self.files.get(log_path, b"") + return {"stdout": data[-int(n):].decode("utf-8", errors="replace"), "stderr": "", "exit_code": 0} return {"stdout": "", "stderr": "", "exit_code": 0} - def upload_file(self, sandbox_id: str, source, destination: str) -> None: + def upload_file(self, external_id: str, source, destination: str) -> None: self.uploads = getattr(self, "uploads", []) - self.uploads.append((sandbox_id, str(source), destination)) + self.uploads.append((external_id, str(source), destination)) - def download_file(self, sandbox_id: str, source: str, destination) -> str: + def download_file(self, external_id: str, source: str, destination, *, user: str | None = None): self.downloads = getattr(self, "downloads", []) - self.downloads.append((sandbox_id, source, str(destination))) - return str(destination) + self.downloads.append((external_id, source, str(destination))) + return destination -def _connection() -> E2BConnectionConfig: - return E2BConnectionConfig( - api_key="super-secret", - api_url="https://e2b.example", - project_id="project-42", - template="doc-compiler", - ) +def _spec() -> dict: + return { + "template": "doc-compiler", + "api_key": "super-secret", + "api_url": "https://e2b.example", + "project_id": "project-42", + "timeout": 600, + } + +def _persisted_spec() -> dict: + return {key: value for key, value in _spec().items() if key != "api_key"} -def test_submit_e2b_persists_sandbox_without_api_key_and_is_idempotent(tmp_path) -> None: - adapter = _FakeE2BAdapter() - service = RemoteJobService(RemoteJobStore(tmp_path / "remote-jobs.db"), e2b_adapter=adapter) - job = service.submit_e2b( +def test_submit_job_persists_sandbox_without_api_key_and_is_idempotent(tmp_path) -> None: + adapter = _FakeAdapter() + service = RemoteJobService(RemoteJobStore(tmp_path / "remote-jobs.db"), adapter_overrides={"e2b": adapter}) + + job = service.submit_job( owner_id="alice", session_id="session-1", + provider="e2b", idempotency_key="session-1:node-1:1", - connection=_connection(), + spec=_spec(), + persisted_specification=_persisted_spec(), ) - replay = service.submit_e2b( + replay = service.submit_job( owner_id="alice", session_id="session-1", + provider="e2b", idempotency_key="session-1:node-1:1", - connection=_connection(), + spec=_spec(), + persisted_specification=_persisted_spec(), ) assert job["status"] == "running" @@ -69,34 +128,128 @@ def test_submit_e2b_persists_sandbox_without_api_key_and_is_idempotent(tmp_path) assert len(adapter.created_specs) == 1 -def test_e2b_job_controls_update_durable_state(tmp_path) -> None: - adapter = _FakeE2BAdapter() - service = RemoteJobService(RemoteJobStore(tmp_path / "remote-jobs.db"), e2b_adapter=adapter) - job = service.submit_e2b( +def test_submit_job_retries_after_a_creation_failure(tmp_path) -> None: + """A job that failed before acquiring an external ID must not poison its + + idempotency key forever: the next submission with the same key resets the + record and re-attempts provider creation.""" + + class _FlakyAdapter(_FakeAdapter): + def __init__(self) -> None: + super().__init__() + self.create_calls = 0 + + def create(self, spec: dict) -> str: + self.create_calls += 1 + if self.create_calls == 1: + raise ValueError("dictionary update sequence element #0 has length 1; 2 is required") + return super().create(spec) + + adapter = _FlakyAdapter() + service = RemoteJobService(RemoteJobStore(tmp_path / "remote-jobs.db"), adapter_overrides={"e2b": adapter}) + + failed = service.submit_job( + owner_id="alice", + session_id="session-1", + provider="e2b", + idempotency_key="session-1:node-1:1", + spec=_spec(), + ) + assert failed["status"] == "failed" + assert failed["external_id"] is None + assert "dictionary update sequence" in failed["error"] + + retried = service.submit_job( + owner_id="alice", + session_id="session-1", + provider="e2b", + idempotency_key="session-1:node-1:1", + spec=_spec(), + ) + assert retried["job_id"] == failed["job_id"] + assert retried["status"] == "running" + assert retried["external_id"] == "sandbox-123" + assert retried["error"] is None + assert adapter.create_calls == 2 + + +def test_submit_job_does_not_retry_a_failure_that_has_an_external_id(tmp_path) -> None: + adapter = _FakeAdapter() + store = RemoteJobStore(tmp_path / "remote-jobs.db") + service = RemoteJobService(store, adapter_overrides={"e2b": adapter}) + job = service.submit_job( owner_id="alice", session_id="session-1", + provider="e2b", idempotency_key="session-1:node-1:1", - connection=_connection(), + spec=_spec(), ) + store.transition_job(job["job_id"], "failed", error="provider died mid-run") - paused = service.pause_e2b(job["job_id"]) + replay = service.submit_job( + owner_id="alice", + session_id="session-1", + provider="e2b", + idempotency_key="session-1:node-1:1", + spec=_spec(), + ) + + assert replay["status"] == "failed" + assert replay["error"] == "provider died mid-run" + assert len(adapter.created_specs) == 1 + + +def test_job_controls_update_durable_state(tmp_path) -> None: + adapter = _FakeAdapter() + service = RemoteJobService(RemoteJobStore(tmp_path / "remote-jobs.db"), adapter_overrides={"e2b": adapter}) + job = service.submit_job( + owner_id="alice", + session_id="session-1", + provider="e2b", + idempotency_key="session-1:node-1:1", + spec=_spec(), + ) + + paused = service.pause_job(job["job_id"]) assert paused["status"] == "paused" assert adapter.paused == ["sandbox-123"] - terminated = service.terminate_e2b(paused["job_id"]) + terminated = service.terminate_job(paused["job_id"]) assert terminated["status"] == "terminated" - assert adapter.terminated == ["sandbox-123"] + assert adapter.cancelled == ["sandbox-123"] -def test_e2b_command_merges_telemetry_after_monitor_observation(tmp_path) -> None: - adapter = _FakeE2BAdapter() +def test_pause_job_raises_capability_error_for_pause_unsupported_provider(tmp_path) -> None: + class _NoPauseAdapter(_FakeAdapter): + capabilities = frozenset({RemoteJobCapability.INTERACTIVE_EXEC, RemoteJobCapability.FILE_TRANSFER}) + + adapter = _NoPauseAdapter() + service = RemoteJobService(RemoteJobStore(tmp_path / "remote-jobs.db"), adapter_overrides={"e2b": adapter}) + job = service.submit_job( + owner_id="alice", + session_id="session-1", + provider="e2b", + idempotency_key="session-1:node-1:1", + spec=_spec(), + ) + + try: + service.pause_job(job["job_id"]) + assert False, "expected CapabilityError" + except Exception as exc: + assert "does not support 'pause'" in str(exc) + + +def test_command_merges_telemetry_after_monitor_observation(tmp_path) -> None: + adapter = _FakeAdapter() store = RemoteJobStore(tmp_path / "remote-jobs.db") - service = RemoteJobService(store, e2b_adapter=adapter) - job = service.submit_e2b( + service = RemoteJobService(store, adapter_overrides={"e2b": adapter}) + job = service.submit_job( owner_id="alice", session_id="session-1", + provider="e2b", idempotency_key="session-1:node-1:1", - connection=_connection(), + spec=_spec(), ) adapter.on_run = lambda: store.record_observation( @@ -105,7 +258,7 @@ def test_e2b_command_merges_telemetry_after_monitor_observation(tmp_path) -> Non expected_revision=store.get_job(job["job_id"])["state_revision"], ) - assert service.run_e2b_command(job["job_id"], "echo done") == { + assert service.run_job_command(job["job_id"], "echo done") == { "stdout": "", "stderr": "", "exit_code": 0 } assert store.get_job(job["job_id"])["snapshot"] == { @@ -115,25 +268,205 @@ def test_e2b_command_merges_telemetry_after_monitor_observation(tmp_path) -> Non } -def test_download_e2b_file_merges_telemetry_and_returns_paths(tmp_path) -> None: - adapter = _FakeE2BAdapter() +def test_download_job_file_merges_telemetry_and_returns_paths(tmp_path) -> None: + adapter = _FakeAdapter() store = RemoteJobStore(tmp_path / "remote-jobs.db") - service = RemoteJobService(store, e2b_adapter=adapter) - job = service.submit_e2b( + service = RemoteJobService(store, adapter_overrides={"e2b": adapter}) + job = service.submit_job( owner_id="alice", session_id="session-1", + provider="e2b", idempotency_key="session-1:node-1:1", - connection=_connection(), + spec=_spec(), ) dest = tmp_path / "CHGCAR" - result = service.download_e2b_file(job["job_id"], "/home/user/CHGCAR", dest) + result = service.download_job_file(job["job_id"], "/home/user/CHGCAR", dest) assert result == {"source": "/home/user/CHGCAR", "destination": str(dest.resolve())} assert adapter.downloads == [("sandbox-123", "/home/user/CHGCAR", str(dest.resolve()))] assert store.get_job(job["job_id"])["snapshot"] == { "provider_status": "reachable", - "sandbox_id": "sandbox-123", "last_download": "CHGCAR", } + +def test_reconcile_job_transitions_on_normalized_status_change(tmp_path) -> None: + class _BatchAdapter(_FakeAdapter): + provider = "bohr_job" + capabilities = frozenset({RemoteJobCapability.BATCH_COLLECT}) + + def __init__(self) -> None: + super().__init__() + self.status_calls = 0 + + def status(self, external_id: str) -> RemoteJobStatus: + # First probe (right after create, inside submit_job) reports + # "queued"; only a later explicit reconcile reports "succeeded" — + # this exercises reconcile_job's own transition, not submission. + self.status_calls += 1 + if self.status_calls == 1: + return RemoteJobStatus(normalized_status="queued", snapshot={"phase": "pending"}) + return RemoteJobStatus(normalized_status="succeeded", snapshot={"phase": "completed"}) + + adapter = _BatchAdapter() + store = RemoteJobStore(tmp_path / "remote-jobs.db") + service = RemoteJobService(store, adapter_overrides={"bohr_job": adapter}) + job = service.submit_job( + owner_id="alice", + session_id="session-1", + provider="bohr_job", + idempotency_key="session-1:node-1:1", + spec={"project_id": 1, "job_name": "n", "machine_type": "c2", "image_address": "img", "command": "cmd"}, + ) + assert job["status"] == "queued" + + reconciled = service.reconcile_job(job["job_id"]) + assert reconciled["status"] == "succeeded" + assert reconciled["snapshot"]["phase"] == "completed" + + +def test_collect_job_outputs_is_idempotent(tmp_path) -> None: + class _BatchAdapter(_FakeAdapter): + provider = "bohr_job" + capabilities = frozenset({RemoteJobCapability.BATCH_COLLECT}) + + def __init__(self) -> None: + super().__init__() + self.collect_calls: list[str] = [] + + def status(self, external_id: str) -> RemoteJobStatus: + return RemoteJobStatus(normalized_status="succeeded", snapshot={"phase": "completed"}) + + def collect_outputs(self, external_id: str, destination_dir): + self.collect_calls.append(external_id) + return [{"source": external_id, "destination": str(destination_dir)}] + + adapter = _BatchAdapter() + store = RemoteJobStore(tmp_path / "remote-jobs.db") + service = RemoteJobService(store, adapter_overrides={"bohr_job": adapter}) + job = service.submit_job( + owner_id="alice", + session_id="session-1", + provider="bohr_job", + idempotency_key="session-1:node-1:1", + spec={"project_id": 1, "job_name": "n", "machine_type": "c2", "image_address": "img", "command": "cmd"}, + ) + service.reconcile_job(job["job_id"]) + + collected = service.collect_job_outputs(job["job_id"], tmp_path / "out") + assert collected["status"] == "collected" + assert len(collected["artifacts"]) == 1 + assert adapter.collect_calls == ["sandbox-123"] + + replay = service.collect_job_outputs(job["job_id"], tmp_path / "out") + assert replay["status"] == "collected" + assert adapter.collect_calls == ["sandbox-123"] + + +def test_start_job_command_persists_handle_with_derived_marker_paths(tmp_path) -> None: + adapter = _FakeAdapter() + service = RemoteJobService(RemoteJobStore(tmp_path / "remote-jobs.db"), adapter_overrides={"e2b": adapter}) + job = service.submit_job( + owner_id="alice", + session_id="session-1", + provider="e2b", + idempotency_key="session-1:node-1:1", + spec=_spec(), + ) + + result = service.start_job_command(job["job_id"], "sleep 300 && echo done") + + assert result["handle"]["log_path"] == f"/tmp/matcreator-cmd-{job['job_id']}.log" + assert result["handle"]["exit_path"] == f"/tmp/matcreator-cmd-{job['job_id']}.exit" + assert adapter.launched_commands == ["sleep 300 && echo done"] + persisted = service.store.get_job(job["job_id"]) + assert persisted["snapshot"]["background_command"]["log_path"] == result["handle"]["log_path"] + + +def test_start_job_command_base64_round_trips_arbitrary_shell_content(tmp_path) -> None: + """Quotes, `$()`, and backticks in the command must survive intact — + proving the wrapper can't be broken out of or reinterpreted.""" + adapter = _FakeAdapter() + service = RemoteJobService(RemoteJobStore(tmp_path / "remote-jobs.db"), adapter_overrides={"e2b": adapter}) + job = service.submit_job( + owner_id="alice", session_id="session-1", provider="e2b", + idempotency_key="session-1:node-1:1", spec=_spec(), + ) + tricky_command = """echo 'it'"'"'s a test' && echo "$(date)" && echo `whoami`""" + + service.start_job_command(job["job_id"], tricky_command) + + assert adapter.launched_commands == [tricky_command] + + +def test_poll_job_command_reports_running_while_no_exit_marker(tmp_path) -> None: + adapter = _FakeAdapter() + service = RemoteJobService(RemoteJobStore(tmp_path / "remote-jobs.db"), adapter_overrides={"e2b": adapter}) + job = service.submit_job( + owner_id="alice", session_id="session-1", provider="e2b", + idempotency_key="session-1:node-1:1", spec=_spec(), + ) + service.start_job_command(job["job_id"], "sleep 300") + + result = service.poll_job_command(job["job_id"]) + + assert result["running"] is True + + +def test_poll_job_command_reports_result_once_finished_and_clears_handle(tmp_path) -> None: + adapter = _FakeAdapter() + service = RemoteJobService(RemoteJobStore(tmp_path / "remote-jobs.db"), adapter_overrides={"e2b": adapter}) + job = service.submit_job( + owner_id="alice", session_id="session-1", provider="e2b", + idempotency_key="session-1:node-1:1", spec=_spec(), + ) + started = service.start_job_command(job["job_id"], "echo hello") + # Simulate the background command finishing inside the sandbox. + adapter.files[started["handle"]["exit_path"]] = b"0\n" + adapter.files[started["handle"]["log_path"]] = b"hello\n" + + result = service.poll_job_command(job["job_id"]) + + assert result == { + "running": False, + "exit_code": 0, + "output_tail": "hello\n", + "log_path": started["handle"]["log_path"], + } + assert service.store.get_job(job["job_id"])["snapshot"]["background_command"] is None + + +def test_poll_job_command_requires_a_started_command(tmp_path) -> None: + adapter = _FakeAdapter() + service = RemoteJobService(RemoteJobStore(tmp_path / "remote-jobs.db"), adapter_overrides={"e2b": adapter}) + job = service.submit_job( + owner_id="alice", session_id="session-1", provider="e2b", + idempotency_key="session-1:node-1:1", spec=_spec(), + ) + + try: + service.poll_job_command(job["job_id"]) + assert False, "expected ValueError" + except ValueError as exc: + assert "no in-flight background command" in str(exc) + + +def test_start_job_command_raises_capability_error_for_batch_provider(tmp_path) -> None: + class _BatchAdapter(_FakeAdapter): + provider = "bohr_job" + capabilities = frozenset({RemoteJobCapability.BATCH_COLLECT}) + + adapter = _BatchAdapter() + service = RemoteJobService(RemoteJobStore(tmp_path / "remote-jobs.db"), adapter_overrides={"bohr_job": adapter}) + job = service.submit_job( + owner_id="alice", session_id="session-1", provider="bohr_job", + idempotency_key="session-1:node-1:1", + spec={"project_id": 1, "job_name": "n", "machine_type": "c2", "image_address": "img", "command": "cmd"}, + ) + + try: + service.start_job_command(job["job_id"], "echo hi") + assert False, "expected CapabilityError" + except Exception as exc: + assert "does not support 'interactive_exec'" in str(exc) diff --git a/tests/test_remote_job_tools.py b/tests/test_remote_job_tools.py new file mode 100644 index 00000000..995bc383 --- /dev/null +++ b/tests/test_remote_job_tools.py @@ -0,0 +1,404 @@ +from __future__ import annotations + +from types import SimpleNamespace + +from matcreator.agents.execution_agent import remote_job_tools +from matcreator.control_plane.providers.e2b import E2BConnectionConfig + + +class _FakeService: + def __init__(self) -> None: + self.submissions: list[dict] = [] + self.store = self + + def submit_job(self, **kwargs): + self.submissions.append(kwargs) + return { + "job_id": "job-123", + "status": "running", + "external_id": "sandbox-123", + } + + def get_job(self, job_id: str): + if job_id != "job-123": + return None + return { + "job_id": job_id, + "owner_id": "alice", + "session_id": "session-1", + "provider": "e2b", + "status": "running", + "external_id": "sandbox-123", + "snapshot": {}, + "error": None, + "updated_at": 1, + } + + def list_events(self, job_id: str): + return [{"event_type": "user_control", "payload": {"action": "terminate", "source": "ui"}}] + + def pause_job(self, job_id: str): + return {"job_id": job_id, "status": "paused", "external_id": "sandbox-123"} + + def terminate_job(self, job_id: str): + return {"job_id": job_id, "status": "terminated", "external_id": "sandbox-123"} + + def run_job_command(self, job_id: str, command: str, *, user: str): + return {"stdout": f"ran {command}", "stderr": "", "exit_code": 0} + + def upload_job_file(self, job_id: str, source, destination: str): + return {"source": str(source), "destination": destination} + + def download_job_file(self, job_id: str, source: str, destination: str): + return {"source": source, "destination": str(destination)} + + def collect_job_outputs(self, job_id: str, destination_dir): + return {"job_id": job_id, "status": "collected", "artifacts": [{"source": "x", "destination": str(destination_dir)}]} + + def start_job_command(self, job_id: str, command: str, *, user: str): + return {"job_id": job_id, "launch": {"stdout": "LAUNCHED\n"}, "handle": {"log_path": "/tmp/x.log", "exit_path": "/tmp/x.exit"}} + + def poll_job_command(self, job_id: str): + return {"running": False, "exit_code": 0, "output_tail": "done\n", "log_path": "/tmp/x.log"} + + +def _context(): + return SimpleNamespace( + state={ + "session_id": "session-1", + "_graph_exec_node_id": "execution_0__node_relax", + "step_number": 2, + }, + _invocation_context=SimpleNamespace(user_id="alice"), + ) + + +def test_submit_e2b_tool_uses_current_session_and_node(monkeypatch) -> None: + service = _FakeService() + monkeypatch.setattr(remote_job_tools, "_service", lambda: service) + monkeypatch.setenv("E2B_API_KEY", "secret") + monkeypatch.setenv("E2B_API_URL", "https://e2b.example") + monkeypatch.setenv("BOHRIUM_PROJECT_ID", "project-42") + + result = remote_job_tools.submit_e2b_sandbox(_context(), timeout=120, template="doc-compiler") + + assert result == { + "status": "running", + "job_id": "job-123", + "sandbox_id": "sandbox-123", + "message": "Tracked E2B sandbox is ready. Use its job_id for status or controls.", + } + submission = service.submissions[0] + assert submission["owner_id"] == "alice" + assert submission["session_id"] == "session-1" + assert submission["provider"] == "e2b" + assert submission["node_id"] == "relax" + assert submission["step_number"] == 2 + assert submission["spec"]["template"] == "doc-compiler" + assert submission["spec"]["api_key"] == "secret" + assert "api_key" not in submission["persisted_specification"] + + +def test_submit_e2b_tool_requires_explicit_template(monkeypatch) -> None: + service = _FakeService() + monkeypatch.setattr(remote_job_tools, "_service", lambda: service) + + result = remote_job_tools.submit_e2b_sandbox(_context()) + + assert result["status"] == "error" + assert "template is required" in result["message"] + assert service.submissions == [] + + +def test_submit_e2b_tool_requires_server_configuration(monkeypatch) -> None: + service = _FakeService() + monkeypatch.setattr(remote_job_tools, "_service", lambda: service) + monkeypatch.delenv("E2B_API_KEY", raising=False) + monkeypatch.setenv("E2B_API_URL", "https://e2b.example") + monkeypatch.delenv("BOHRIUM_PROJECT_ID", raising=False) + + result = remote_job_tools.submit_e2b_sandbox(_context(), template="doc-compiler") + + assert result["status"] == "error" + assert "E2B_API_KEY" in result["message"] + assert "BOHRIUM_PROJECT_ID" in result["message"] + assert "E2B_API_URL" not in result["message"] + assert service.submissions == [] + + +def test_submit_e2b_tool_coerces_json_string_lifecycle(monkeypatch) -> None: + service = _FakeService() + monkeypatch.setattr(remote_job_tools, "_service", lambda: service) + monkeypatch.setenv("E2B_API_KEY", "secret") + monkeypatch.setenv("E2B_API_URL", "https://e2b.example") + monkeypatch.setenv("BOHRIUM_PROJECT_ID", "project-42") + + result = remote_job_tools.submit_e2b_sandbox( + _context(), template="doc-compiler", lifecycle='{"on_timeout": "pause", "auto_resume": false}' + ) + + assert result["status"] == "running" + assert service.submissions[0]["spec"]["lifecycle"] == {"on_timeout": "pause", "auto_resume": False} + + +def test_submit_e2b_tool_rejects_non_object_lifecycle(monkeypatch) -> None: + service = _FakeService() + monkeypatch.setattr(remote_job_tools, "_service", lambda: service) + monkeypatch.setenv("E2B_API_KEY", "secret") + monkeypatch.setenv("E2B_API_URL", "https://e2b.example") + monkeypatch.setenv("BOHRIUM_PROJECT_ID", "project-42") + + for bad_lifecycle in ("pause on timeout", '["pause"]'): + result = remote_job_tools.submit_e2b_sandbox( + _context(), template="doc-compiler", lifecycle=bad_lifecycle + ) + assert result["status"] == "error" + assert "lifecycle must be a JSON object" in result["message"] + assert service.submissions == [] + + +def test_submit_e2b_tool_surfaces_failed_job_error_instead_of_claiming_ready(monkeypatch) -> None: + service = _FakeService() + + def _failed_submit(**kwargs): + service.submissions.append(kwargs) + return { + "job_id": "job-123", + "status": "failed", + "external_id": None, + "error": "e2b job creation failed: boom", + } + + service.submit_job = _failed_submit + monkeypatch.setattr(remote_job_tools, "_service", lambda: service) + monkeypatch.setenv("E2B_API_KEY", "secret") + monkeypatch.setenv("E2B_API_URL", "https://e2b.example") + monkeypatch.setenv("BOHRIUM_PROJECT_ID", "project-42") + + result = remote_job_tools.submit_e2b_sandbox(_context(), template="doc-compiler") + + assert result["status"] == "failed" + assert result["sandbox_id"] is None + assert "e2b job creation failed: boom" in result["message"] + assert "ready" not in result["message"] + + +def test_e2b_connection_uses_configured_environment_names(monkeypatch) -> None: + monkeypatch.setenv("E2B_API_KEY", "access-key") + monkeypatch.setenv("E2B_API_URL", "https://e2b.example") + monkeypatch.setenv("BOHRIUM_PROJECT_ID", "project-7") + + connection = remote_job_tools._connection() + + assert connection == E2BConnectionConfig( + api_key="access-key", + api_url="https://e2b.example", + project_id="project-7", + template="", + ) + + +def test_submit_bohr_sandbox_tool_requires_project_id(monkeypatch) -> None: + service = _FakeService() + monkeypatch.setattr(remote_job_tools, "_service", lambda: service) + monkeypatch.delenv("BOHRIUM_PROJECT_ID", raising=False) + + result = remote_job_tools.submit_bohr_sandbox(_context()) + + assert result["status"] == "error" + assert "project_id" in result["message"] + assert service.submissions == [] + + +def test_submit_bohr_sandbox_tool_requires_explicit_template(monkeypatch) -> None: + service = _FakeService() + monkeypatch.setattr(remote_job_tools, "_service", lambda: service) + + result = remote_job_tools.submit_bohr_sandbox(_context(), project_id=42) + + assert result["status"] == "error" + assert "template is required" in result["message"] + assert service.submissions == [] + + +def test_submit_bohr_sandbox_tool_submits_with_provider(monkeypatch) -> None: + service = _FakeService() + monkeypatch.setattr(remote_job_tools, "_service", lambda: service) + + result = remote_job_tools.submit_bohr_sandbox(_context(), project_id=42, template="sdbxagent") + + assert result["status"] == "running" + assert result["sandbox_id"] == "sandbox-123" + assert service.submissions[0]["provider"] == "bohr_sandbox" + assert service.submissions[0]["spec"]["project_id"] == 42 + + +def test_submit_bohr_job_tool_requires_all_fields(monkeypatch) -> None: + service = _FakeService() + monkeypatch.setattr(remote_job_tools, "_service", lambda: service) + monkeypatch.delenv("BOHRIUM_PROJECT_ID", raising=False) + + result = remote_job_tools.submit_bohr_job(_context(), job_name="n") + + assert result["status"] == "error" + assert "project_id" in result["message"] + assert "machine_type" in result["message"] + assert service.submissions == [] + + +def test_submit_bohr_job_tool_submits_batch_spec(monkeypatch) -> None: + service = _FakeService() + monkeypatch.setattr(remote_job_tools, "_service", lambda: service) + + result = remote_job_tools.submit_bohr_job( + _context(), + project_id=42, + job_name="relax-job", + machine_type="c8_m32_cpu", + image_address="registry.dp.tech/dptech/vasp:5.4.4", + command="vasp_std", + ) + + assert result["status"] == "running" + assert result["bohr_job_id"] == "sandbox-123" + submission = service.submissions[0] + assert submission["provider"] == "bohr_job" + assert submission["spec"]["command"] == "vasp_std" + + +def test_remote_job_tools_reject_jobs_from_another_session(monkeypatch) -> None: + service = _FakeService() + monkeypatch.setattr(remote_job_tools, "_service", lambda: service) + context = _context() + context._invocation_context.user_id = "bob" + + assert remote_job_tools.get_remote_job_status("job-123", context) == { + "status": "error", + "message": "Remote job was not found in this session.", + } + + +def test_remote_job_status_exposes_user_control(monkeypatch) -> None: + service = _FakeService() + monkeypatch.setattr(remote_job_tools, "_service", lambda: service) + + status = remote_job_tools.get_remote_job_status("job-123", _context()) + + assert status["user_control"] == {"action": "terminate", "source": "ui"} + + +def test_remote_job_command_and_workspace_upload_are_scoped_to_owned_job(tmp_path, monkeypatch) -> None: + service = _FakeService() + monkeypatch.setattr(remote_job_tools, "_service", lambda: service) + context = _context() + context.state["workspace_dir"] = str(tmp_path) + source = tmp_path / "input.txt" + source.write_text("input", encoding="utf-8") + + assert remote_job_tools.run_remote_job_command("job-123", "echo hello", context) == { + "stdout": "ran echo hello", "stderr": "", "exit_code": 0 + } + assert remote_job_tools.upload_remote_job_input("job-123", "input.txt", "/home/user/input.txt", context) == { + "source": str(source), "destination": "/home/user/input.txt" + } + assert remote_job_tools.upload_remote_job_input( + "job-123", "/tmp/outside.txt", "/tmp/outside.txt", context + )["status"] == "error" + + +def test_download_remote_job_output_is_scoped_to_workspace(tmp_path, monkeypatch) -> None: + service = _FakeService() + monkeypatch.setattr(remote_job_tools, "_service", lambda: service) + context = _context() + context.state["workspace_dir"] = str(tmp_path) + destination = tmp_path / "outputs" / "CHGCAR" + + result = remote_job_tools.download_remote_job_output( + "job-123", "/home/user/CHGCAR", str(destination), context + ) + assert result == {"source": "/home/user/CHGCAR", "destination": str(destination.resolve())} + + assert remote_job_tools.download_remote_job_output( + "job-123", "/home/user/CHGCAR", "/tmp/outside.txt", context + )["status"] == "error" + + +def test_download_remote_job_output_rejects_missing_workspace_dir(monkeypatch) -> None: + service = _FakeService() + monkeypatch.setattr(remote_job_tools, "_service", lambda: service) + context = _context() # no workspace_dir set + + result = remote_job_tools.download_remote_job_output( + "job-123", "/home/user/CHGCAR", "outputs/CHGCAR", context + ) + assert result["status"] == "error" + assert "workspace_dir" in result["message"] + + +def test_collect_remote_job_outputs_is_scoped_to_workspace(tmp_path, monkeypatch) -> None: + service = _FakeService() + monkeypatch.setattr(remote_job_tools, "_service", lambda: service) + context = _context() + context.state["workspace_dir"] = str(tmp_path) + + result = remote_job_tools.collect_remote_job_outputs("job-123", "outputs", context) + + assert result["status"] == "collected" + assert result["artifacts"][0]["destination"] == str((tmp_path / "outputs").resolve()) + + assert remote_job_tools.collect_remote_job_outputs( + "job-123", "/tmp/outside", context + )["status"] == "error" + + +def test_start_remote_job_command_delegates_to_service(monkeypatch) -> None: + service = _FakeService() + monkeypatch.setattr(remote_job_tools, "_service", lambda: service) + + result = remote_job_tools.start_remote_job_command("job-123", "sleep 300", _context()) + + assert result["job_id"] == "job-123" + assert result["handle"]["log_path"] == "/tmp/x.log" + + +def test_start_remote_job_command_rejects_jobs_from_another_session(monkeypatch) -> None: + service = _FakeService() + monkeypatch.setattr(remote_job_tools, "_service", lambda: service) + context = _context() + context._invocation_context.user_id = "bob" + + result = remote_job_tools.start_remote_job_command("job-123", "sleep 300", context) + + assert result == { + "status": "error", + "message": "Remote job was not found in this session.", + } + + +def test_poll_remote_job_command_delegates_to_service(monkeypatch) -> None: + service = _FakeService() + monkeypatch.setattr(remote_job_tools, "_service", lambda: service) + + result = remote_job_tools.poll_remote_job_command("job-123", _context()) + + assert result == { + "running": False, + "exit_code": 0, + "output_tail": "done\n", + "log_path": "/tmp/x.log", + } + + +def test_poll_remote_job_command_surfaces_service_errors(monkeypatch) -> None: + service = _FakeService() + + def _boom(job_id): + raise ValueError("no in-flight background command") + + service.poll_job_command = _boom + monkeypatch.setattr(remote_job_tools, "_service", lambda: service) + + result = remote_job_tools.poll_remote_job_command("job-123", _context()) + + assert result["status"] == "error" + assert "no in-flight background command" in result["message"] diff --git a/tests/test_step_executor_runner.py b/tests/test_step_executor_runner.py index 340b574e..a78e9f57 100644 --- a/tests/test_step_executor_runner.py +++ b/tests/test_step_executor_runner.py @@ -19,18 +19,23 @@ from matcreator.agents.session_log import SESSION_ARTIFACTS_KEY -def test_step_executor_registers_tracked_e2b_tools() -> None: +def test_step_executor_registers_tracked_remote_job_tools() -> None: agent = build_step_executor_agent(LLMCard(name="test", model="test-model")) tool_names = {tool.name for tool in agent.tools if hasattr(tool, "name")} assert { "submit_e2b_sandbox", - "get_e2b_job_status", - "run_e2b_command", - "upload_e2b_input", - "download_e2b_output", - "pause_e2b_sandbox", - "terminate_e2b_sandbox", + "submit_bohr_sandbox", + "submit_bohr_job", + "get_remote_job_status", + "run_remote_job_command", + "start_remote_job_command", + "poll_remote_job_command", + "upload_remote_job_input", + "download_remote_job_output", + "collect_remote_job_outputs", + "pause_remote_job", + "terminate_remote_job", } <= tool_names @@ -431,7 +436,7 @@ def test_resumed_node_receives_explicit_reattachment_instructions(): assert context is not None assert "job-1" in context assert "sbx-1" in context - assert "get_e2b_job_status" in context + assert "get_remote_job_status" in context assert "Do NOT call submit_e2b_sandbox" in context diff --git a/web/main.py b/web/main.py index f0b1e549..26755cf3 100644 --- a/web/main.py +++ b/web/main.py @@ -99,6 +99,7 @@ from matcreator.control_plane.remote_job_monitor import RemoteJobMonitor # noqa: E402 from matcreator.control_plane.remote_job_service import RemoteJobService # noqa: E402 from matcreator.control_plane.remote_jobs import RemoteJobStore # noqa: E402 +from matcreator.control_plane.providers import CapabilityError # noqa: E402 from matcreator.control_plane.benchmark_client import BenchmarkApiError, BenchmarkClient, sanitize_bank_id # noqa: E402 from matcreator.control_plane.evaluation_manager import EvaluationManager # noqa: E402 from matcreator.control_plane.evaluation_runtime import RuntimeOutcome, RuntimeSpec # noqa: E402 @@ -2669,12 +2670,14 @@ async def pause_session_remote_job( job_id: str, user_id: str = Query(..., description="Current signed-in user"), ) -> JSONResponse: - """Pause one E2B sandbox and notify its linked executor without stopping it.""" + """Pause one remote job (if its provider supports pausing) and notify its linked executor without stopping it.""" job = _get_owned_remote_job(session_id, job_id, user_id) try: - paused = await asyncio.to_thread(_remote_job_service_for_owner(user_id).pause_e2b, job_id) + paused = await asyncio.to_thread(_remote_job_service_for_owner(user_id).pause_job, job_id) except (KeyError, ValueError) as exc: raise HTTPException(status_code=409, detail=str(exc)) from exc + except CapabilityError as exc: + raise HTTPException(status_code=409, detail=str(exc)) from exc await asyncio.to_thread( _remote_job_store_for_owner(user_id).record_user_control, job_id, @@ -2689,10 +2692,10 @@ async def terminate_session_remote_job( job_id: str, user_id: str = Query(..., description="Current signed-in user"), ) -> JSONResponse: - """Terminate one E2B sandbox and notify its linked executor without stopping it.""" + """Terminate one remote job and notify its linked executor without stopping it.""" job = _get_owned_remote_job(session_id, job_id, user_id) try: - terminated = await asyncio.to_thread(_remote_job_service_for_owner(user_id).terminate_e2b, job_id) + terminated = await asyncio.to_thread(_remote_job_service_for_owner(user_id).terminate_job, job_id) except (KeyError, ValueError) as exc: raise HTTPException(status_code=409, detail=str(exc)) from exc await asyncio.to_thread( @@ -2709,12 +2712,10 @@ async def refresh_session_remote_job( job_id: str, user_id: str = Query(..., description="Current signed-in user"), ) -> JSONResponse: - """Synchronize a caller-owned active E2B job with its sandbox.""" + """Synchronize a caller-owned active remote job with its provider.""" job = _get_owned_remote_job(session_id, job_id, user_id) - if job["provider"] != "e2b": - raise HTTPException(status_code=409, detail="Remote job is not managed by E2B") try: - refreshed = await asyncio.to_thread(_remote_job_service_for_owner(user_id).reconcile_e2b, job_id) + refreshed = await asyncio.to_thread(_remote_job_service_for_owner(user_id).reconcile_job, job_id) except (KeyError, ValueError) as exc: raise HTTPException(status_code=409, detail=str(exc)) from exc return JSONResponse(refreshed) @@ -3770,7 +3771,7 @@ async def cancel_session_execution( paused_jobs = [] if user_id: paused_jobs = await asyncio.to_thread( - _remote_job_service_for_owner(user_id).pause_active_session_e2b_jobs, + _remote_job_service_for_owner(user_id).pause_active_session_jobs, owner_id=user_id, session_id=session_id, )