refactor(schedule): extract execution/ + runners in layered refactor (stage 4)
- Move worker.py -> execution/worker.py, executor.py -> execution/executor.py (byte-identical copies; import sites updated) - Merge old execution.py + notebook_runner.py into execution/runners/notebook.py: subprocess CLI (main/emit_outputs) plus the in-process helpers (_execute_notebook/_execute_python/execute_artifact) - schedule/notebook_runner.py becomes a compatibility shim so `python -m schedule.notebook_runner` (the worker's stable -m string) still works - Delete flat execution.py (shadowed by the new execution/ package) - Zero behavior change; schedule/pyproject.toml untouched Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
tao.chen
@@ -0,0 +1,4 @@
|
||||
"""
|
||||
@Time :2026/7/29
|
||||
@Author :tao.chen
|
||||
"""
|
||||
+107
@@ -1,12 +1,29 @@
|
||||
"""Notebook / python-script execution backend for the schedule worker.
|
||||
|
||||
Two responsibilities in one module:
|
||||
|
||||
- Subprocess CLI entry, invoked as ``python -m schedule.notebook_runner``
|
||||
through the shim at ``schedule/notebook_runner.py``: executes a notebook
|
||||
out-of-process with nbclient and writes the executed artifact.
|
||||
- In-process runner helpers used by :func:`execute_artifact`: resolve the
|
||||
target Python binary, bound the console log to ``MAX_LOG_BYTES``, and run
|
||||
a notebook or a plain python script with a wall-clock timeout.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import sys
|
||||
import tempfile
|
||||
import traceback
|
||||
from pathlib import Path, PurePosixPath
|
||||
|
||||
import nbformat
|
||||
from loguru import logger
|
||||
from nbclient import NotebookClient
|
||||
|
||||
from schedule.domain.execution import ExecutionResult
|
||||
|
||||
MAX_LOG_BYTES = 4 * 1024 * 1024
|
||||
@@ -261,3 +278,93 @@ async def execute_artifact(
|
||||
)
|
||||
logger.error("unsupported script_type: {}", script_type)
|
||||
raise ValueError(f"unsupported script_type: {script_type}")
|
||||
|
||||
|
||||
def emit_outputs(notebook: object) -> None:
|
||||
for cell in notebook.cells: # type: ignore[attr-defined]
|
||||
if cell.get("cell_type") != "code":
|
||||
continue
|
||||
for output in cell.get("outputs", []):
|
||||
output_type = output.get("output_type")
|
||||
if output_type == "stream":
|
||||
text = output.get("text", "")
|
||||
print(
|
||||
"".join(text) if isinstance(text, list) else str(text),
|
||||
end="",
|
||||
flush=True,
|
||||
)
|
||||
elif output_type == "error":
|
||||
print(
|
||||
f"{output.get('ename', 'Error')}: {output.get('evalue', '')}",
|
||||
file=sys.stderr,
|
||||
flush=True,
|
||||
)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--input", required=True)
|
||||
parser.add_argument("--output", required=True)
|
||||
parser.add_argument("--timeout", required=True, type=int)
|
||||
parser.add_argument(
|
||||
"--python-version",
|
||||
choices=("3.8", "3.10", "3.12"),
|
||||
default="3.12",
|
||||
)
|
||||
parser.add_argument("--arguments-json", default="[]")
|
||||
args = parser.parse_args()
|
||||
|
||||
source = Path(args.input)
|
||||
output = Path(args.output)
|
||||
arguments = json.loads(args.arguments_json)
|
||||
if not isinstance(arguments, list) or not all(
|
||||
isinstance(item, str) for item in arguments
|
||||
):
|
||||
raise ValueError("arguments-json must contain an array of strings")
|
||||
|
||||
logger.info(
|
||||
"notebook runner start: input={} timeout={}s python={} args={}",
|
||||
source.name,
|
||||
args.timeout,
|
||||
args.python_version,
|
||||
len(arguments),
|
||||
)
|
||||
notebook = nbformat.read(source, as_version=4)
|
||||
if arguments:
|
||||
notebook.cells.insert(
|
||||
0,
|
||||
nbformat.v4.new_code_cell(
|
||||
"import sys\n"
|
||||
f"sys.argv = {json.dumps([source.name, *arguments], ensure_ascii=False)}",
|
||||
metadata={"tags": ["injected-parameters"]},
|
||||
),
|
||||
)
|
||||
exit_code = 0
|
||||
try:
|
||||
kernel_name = f"python{args.python_version.replace('.', '')}"
|
||||
client = NotebookClient(
|
||||
notebook,
|
||||
timeout=max(1, args.timeout),
|
||||
kernel_name=kernel_name,
|
||||
allow_errors=False,
|
||||
)
|
||||
logger.debug(
|
||||
"notebook client created: kernel={} timeout={}s",
|
||||
kernel_name,
|
||||
max(1, args.timeout),
|
||||
)
|
||||
# No explicit cwd — the kernel inherits the parent's cwd, which the
|
||||
# scheduler sets to the staged artifact directory. Keeping it here
|
||||
# avoids any "cwd must exist" requirement on the host.
|
||||
client.execute()
|
||||
logger.info("notebook client execute done: input={}", source.name)
|
||||
except Exception as exc:
|
||||
logger.exception("notebook execute failed: input={}", source.name)
|
||||
traceback.print_exc()
|
||||
exit_code = 124 if "timeout" in type(exc).__name__.lower() else 1
|
||||
finally:
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
nbformat.write(notebook, output)
|
||||
emit_outputs(notebook)
|
||||
logger.debug("notebook output written: {}", output)
|
||||
raise SystemExit(exit_code)
|
||||
@@ -0,0 +1,556 @@
|
||||
"""Node-level worker: executes one ``job.node.execute`` event.
|
||||
|
||||
The orchestrator (see ``schedule.orchestrator``) writes a ``job.node.execute``
|
||||
Outbox row with all the metadata needed to run the node (script type,
|
||||
artifact location, timeout, arguments ...). The polling loop picks those up
|
||||
and calls :meth:`NodeExecutor.handle_node_execute`. This module owns the
|
||||
artifact download + subprocess invocation + result-upload side of things.
|
||||
|
||||
`DispatchOrchestrator` writes the node's lifecycle row + outbox event; the
|
||||
worker only mutates ``ScheduleNodeRuns`` columns related to execution
|
||||
(started_at / finished_at / exit_code / result_object_id ...).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import traceback
|
||||
from typing import Any
|
||||
|
||||
from common.config import settings
|
||||
from common.db import session_scope
|
||||
from common.db.models import (
|
||||
ConsumerInbox,
|
||||
ScheduleNodeRuns,
|
||||
ScheduleNodes,
|
||||
ScheduleRuns,
|
||||
Schedules,
|
||||
StorageObjects,
|
||||
Users,
|
||||
Versions,
|
||||
Workspaces,
|
||||
)
|
||||
from common.eventing import (
|
||||
add_outbox_event,
|
||||
event_time,
|
||||
schedule_event_type,
|
||||
utcnow,
|
||||
)
|
||||
from common.scheduler.trigger import SYSTEM_CRON_USER_ID
|
||||
from loguru import logger
|
||||
from sqlalchemy import select
|
||||
|
||||
from schedule.domain.context import TERMINAL_NODE_STATES
|
||||
from schedule.domain.execution import ExecutionResult
|
||||
from schedule.execution.runners.notebook import execute_artifact
|
||||
|
||||
NODE_EXECUTE_EVENT = schedule_event_type("job.node.execute")
|
||||
NODE_FINISHED_EVENT = schedule_event_type("job.node.finished")
|
||||
|
||||
class NodeExecutor:
|
||||
"""Owns the actual execution of one schedule node (notebook / python)."""
|
||||
|
||||
# P0-5 / C1: distinct error_code for runs blocked because the originating
|
||||
# user was disabled or soft-deleted between queue time and worker pickup.
|
||||
# The value lands in the NODE_FINISHED_EVENT outbox payload (the
|
||||
# ``error_code`` field) — schedule_node_runs has no such column; the row
|
||||
# only carries the message text. Operators grep the outbox stream.
|
||||
USER_DISABLED_ERROR_CODE = "USER_DISABLED"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
session_factory,
|
||||
object_store: Any,
|
||||
storage_client: Any,
|
||||
) -> None:
|
||||
self.session_factory = session_factory
|
||||
self.object_store = object_store
|
||||
self.storage_client = storage_client
|
||||
# Per-bucket object store cache. The injected ``object_store`` is
|
||||
# the default version-bucket store; workspaces that override
|
||||
# ``Workspaces.artifact_bucket`` need a store bound to that custom
|
||||
# bucket (P0-3 fix). Build lazily so the common (no-override) path
|
||||
# incurs no extra cost.
|
||||
self._bucket_stores: dict[str, Any] = {
|
||||
settings.s3_version_bucket: object_store,
|
||||
}
|
||||
|
||||
def _store_for(self, bucket_name: str) -> Any:
|
||||
"""Return the AsyncStorageBackend bound to ``bucket_name``.
|
||||
|
||||
Caches per-bucket stores on first use; the default version bucket
|
||||
always reuses the injected ``object_store`` so the common path
|
||||
stays zero-allocation.
|
||||
"""
|
||||
store = self._bucket_stores.get(bucket_name)
|
||||
if store is not None:
|
||||
return store
|
||||
from schedule.service import build_object_store
|
||||
|
||||
store = build_object_store(bucket_name=bucket_name)
|
||||
self._bucket_stores[bucket_name] = store
|
||||
return store
|
||||
|
||||
async def handle_node_execute(
|
||||
self,
|
||||
event: dict[str, Any],
|
||||
message_id: str,
|
||||
) -> None:
|
||||
if event.get("event_type") != NODE_EXECUTE_EVENT:
|
||||
raise ValueError("unexpected event type")
|
||||
payload = event["payload"]
|
||||
if not await self._set_node_running(event, message_id):
|
||||
return
|
||||
logger.info(
|
||||
"node execute start: node_run={} script_type={} timeout={}s",
|
||||
payload["node_run_id"][-12:],
|
||||
payload["script_type"],
|
||||
payload["timeout_seconds"],
|
||||
)
|
||||
started_at = utcnow()
|
||||
context: dict[str, Any] | None = None
|
||||
try:
|
||||
context = await self._execution_context(payload)
|
||||
content = await self._download_artifact(
|
||||
bucket_name=context["bucket_name"],
|
||||
object_key=context["object_key"],
|
||||
content_hash=context["content_hash"],
|
||||
)
|
||||
python_version = await self._node_python_version(
|
||||
payload["node_run_id"]
|
||||
)
|
||||
result = await execute_artifact(
|
||||
content,
|
||||
run_id=payload["run_id"],
|
||||
node_run_id=payload["node_run_id"],
|
||||
script_type=payload["script_type"],
|
||||
artifact_path=payload["artifact_path"],
|
||||
arguments=[str(item) for item in payload.get("arguments", [])],
|
||||
timeout_seconds=int(payload["timeout_seconds"]),
|
||||
python_version=python_version,
|
||||
)
|
||||
except Exception as exc:
|
||||
trace = traceback.format_exc()
|
||||
# P0-5 / C1: 让 _assert_user_active 抛的 ValueError 透传成单独的
|
||||
# error_code,便于运维 grep 区分"用户被禁用"和"代码崩溃"。
|
||||
exc_message = str(exc)
|
||||
error_code = (
|
||||
self.USER_DISABLED_ERROR_CODE
|
||||
if exc_message.startswith("USER_DISABLED:")
|
||||
else "WORKER_EXECUTION_FAILED"
|
||||
)
|
||||
result = ExecutionResult(
|
||||
status="failed",
|
||||
exit_code=1,
|
||||
logs=trace.encode("utf-8", errors="replace"),
|
||||
result=json.dumps(
|
||||
{"status": "failed", "error": exc_message},
|
||||
ensure_ascii=False,
|
||||
).encode("utf-8"),
|
||||
result_file_name=f"{payload['node_run_id']}-result.json",
|
||||
result_content_type="application/json",
|
||||
error_code=error_code,
|
||||
error_message=exc_message[:2000],
|
||||
)
|
||||
if context is None:
|
||||
context = await self._fallback_execution_context(payload)
|
||||
|
||||
log_id, result_id, upload_error = await self._upload_execution_artifacts(
|
||||
payload=payload,
|
||||
context=context,
|
||||
result=result,
|
||||
)
|
||||
finished_at = utcnow()
|
||||
duration_ms = max(
|
||||
0,
|
||||
int((finished_at - started_at).total_seconds() * 1000),
|
||||
)
|
||||
error_message = result.error_message
|
||||
if upload_error:
|
||||
error_message = (
|
||||
f"{error_message}; {upload_error}"
|
||||
if error_message
|
||||
else upload_error
|
||||
)[:2000]
|
||||
final_status = "failed" if upload_error else result.status
|
||||
final_error_code = (
|
||||
"ARTIFACT_UPLOAD_FAILED" if upload_error else result.error_code
|
||||
)
|
||||
|
||||
async with session_scope(self.session_factory) as session:
|
||||
node_run = await session.scalar(
|
||||
select(ScheduleNodeRuns)
|
||||
.where(
|
||||
ScheduleNodeRuns.node_run_id == payload["node_run_id"],
|
||||
)
|
||||
.with_for_update()
|
||||
)
|
||||
if node_run is None:
|
||||
raise ValueError("schedule node run disappeared")
|
||||
inbox = await session.get(
|
||||
ConsumerInbox,
|
||||
("job-workers", event["event_id"]),
|
||||
with_for_update=True,
|
||||
)
|
||||
if inbox is None:
|
||||
raise ValueError("job worker inbox record disappeared")
|
||||
if node_run.node_status not in TERMINAL_NODE_STATES:
|
||||
node_run.node_status = final_status
|
||||
node_run.finished_at = finished_at
|
||||
node_run.duration_ms = duration_ms
|
||||
node_run.exit_code = result.exit_code
|
||||
node_run.message = (
|
||||
"节点执行成功"
|
||||
if final_status == "succeeded"
|
||||
else (error_message or "节点执行失败")
|
||||
)[:2000]
|
||||
node_run.metrics_json = {
|
||||
"log_size_bytes": len(result.logs),
|
||||
"result_size_bytes": len(result.result),
|
||||
}
|
||||
node_run.logs_object_id = log_id
|
||||
node_run.result_object_id = result_id
|
||||
node_run.state_version += 1
|
||||
await add_outbox_event(
|
||||
session,
|
||||
event_type=NODE_FINISHED_EVENT,
|
||||
producer="job-worker",
|
||||
trace_id=event["trace_id"],
|
||||
aggregate_type="schedule_node_run",
|
||||
aggregate_id=node_run.node_run_id,
|
||||
idempotency_key=(
|
||||
f"{node_run.node_run_id}:{node_run.attempt_no}:finished"
|
||||
),
|
||||
payload={
|
||||
"workspace_id": context["workspace_id"],
|
||||
"run_id": node_run.run_id,
|
||||
"node_run_id": node_run.node_run_id,
|
||||
"node_id": node_run.node_id,
|
||||
"versions_id": node_run.versions_id,
|
||||
"attempt_no": node_run.attempt_no,
|
||||
"node_status": node_run.node_status,
|
||||
"exit_code": node_run.exit_code,
|
||||
"started_at": event_time(
|
||||
node_run.started_at or started_at
|
||||
),
|
||||
"finished_at": event_time(finished_at),
|
||||
"duration_ms": duration_ms,
|
||||
"logs_object_id": log_id,
|
||||
"result_object_id": result_id,
|
||||
"error_code": final_error_code,
|
||||
"error_message": error_message,
|
||||
},
|
||||
)
|
||||
self._finish_inbox(inbox)
|
||||
logger.info(
|
||||
"node execute done: node_run={} status={} error_code={}",
|
||||
payload["node_run_id"][-12:],
|
||||
final_status,
|
||||
final_error_code,
|
||||
)
|
||||
|
||||
async def _set_node_running(
|
||||
self,
|
||||
event: dict[str, Any],
|
||||
message_id: str,
|
||||
) -> bool:
|
||||
payload = event["payload"]
|
||||
async with session_scope(self.session_factory) as session:
|
||||
inbox, should_process = await self._start_inbox(
|
||||
session,
|
||||
consumer_name="job-workers",
|
||||
event_id=event["event_id"],
|
||||
message_id=message_id,
|
||||
)
|
||||
if not should_process:
|
||||
return False
|
||||
node_run = await session.scalar(
|
||||
select(ScheduleNodeRuns)
|
||||
.where(
|
||||
ScheduleNodeRuns.node_run_id == payload["node_run_id"],
|
||||
)
|
||||
.with_for_update()
|
||||
)
|
||||
if node_run is None:
|
||||
raise ValueError("schedule node run does not exist")
|
||||
if node_run.node_status in TERMINAL_NODE_STATES:
|
||||
logger.debug(
|
||||
"node already terminal: node_run={} status={}",
|
||||
payload["node_run_id"][-12:],
|
||||
node_run.node_status,
|
||||
)
|
||||
self._finish_inbox(inbox)
|
||||
return False
|
||||
if node_run.node_status == "queued":
|
||||
node_run.node_status = "running"
|
||||
node_run.started_at = utcnow()
|
||||
node_run.message = "Worker 正在执行稳定版本"
|
||||
node_run.state_version += 1
|
||||
logger.debug(
|
||||
"node running: node_run={}",
|
||||
payload["node_run_id"][-12:],
|
||||
)
|
||||
return True
|
||||
|
||||
async def _node_python_version(
|
||||
self,
|
||||
node_run_id: str,
|
||||
) -> str:
|
||||
async with self.session_factory() as session:
|
||||
row = (
|
||||
await session.execute(
|
||||
select(ScheduleNodes.python_version)
|
||||
.join(
|
||||
ScheduleNodeRuns,
|
||||
ScheduleNodeRuns.node_id == ScheduleNodes.node_id,
|
||||
)
|
||||
.where(ScheduleNodeRuns.node_run_id == node_run_id)
|
||||
)
|
||||
).one_or_none()
|
||||
if row is None:
|
||||
logger.warning(
|
||||
"node python_version not found, defaulting to 3.12: node_run={}",
|
||||
node_run_id[-12:],
|
||||
)
|
||||
return "3.12"
|
||||
return row[0]
|
||||
|
||||
async def _assert_user_active(
|
||||
self,
|
||||
session: AsyncSession,
|
||||
user_id: str,
|
||||
) -> None:
|
||||
"""P0-5 / C1: re-verify the user is still ``status='active'`` and
|
||||
``is_deleted=0`` before executing a run they originated.
|
||||
|
||||
Raises :class:`ValueError` whose message starts with
|
||||
``USER_DISABLED:`` when the user has been disabled or soft-deleted
|
||||
between run creation and worker pickup. The outer
|
||||
:meth:`handle_node_execute` parses that prefix and routes the
|
||||
resulting ``error_code="USER_DISABLED"`` into the
|
||||
``NODE_FINISHED_EVENT`` outbox payload (the
|
||||
``schedule_node_runs`` row has no ``error_code`` column, only a
|
||||
``message`` text field).
|
||||
"""
|
||||
user = await session.scalar(
|
||||
select(Users.status, Users.is_deleted).where(Users.user_id == user_id)
|
||||
)
|
||||
if user is None:
|
||||
raise ValueError(
|
||||
f"USER_DISABLED: originating user {user_id} no longer exists"
|
||||
)
|
||||
status_value, is_deleted = user
|
||||
if status_value != "active" or is_deleted != 0:
|
||||
raise ValueError(
|
||||
f"USER_DISABLED: originating user {user_id} is "
|
||||
f"status={status_value!r} is_deleted={is_deleted}"
|
||||
)
|
||||
|
||||
async def _execution_context(
|
||||
self,
|
||||
payload: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
async with self.session_factory() as session:
|
||||
row = (
|
||||
await session.execute(
|
||||
select(
|
||||
ScheduleNodeRuns,
|
||||
ScheduleRuns,
|
||||
Versions,
|
||||
StorageObjects,
|
||||
Workspaces,
|
||||
Schedules,
|
||||
)
|
||||
.join(
|
||||
ScheduleRuns,
|
||||
ScheduleRuns.run_id == ScheduleNodeRuns.run_id,
|
||||
)
|
||||
.join(
|
||||
Versions,
|
||||
Versions.versions_id == ScheduleNodeRuns.versions_id,
|
||||
)
|
||||
.join(
|
||||
StorageObjects,
|
||||
StorageObjects.storage_object_id
|
||||
== Versions.artifact_object_id,
|
||||
)
|
||||
.join(
|
||||
Workspaces,
|
||||
Workspaces.workspace_id == ScheduleRuns.workspace_id,
|
||||
)
|
||||
.join(
|
||||
Schedules,
|
||||
Schedules.schedule_id == ScheduleRuns.schedule_id,
|
||||
)
|
||||
.where(
|
||||
ScheduleNodeRuns.node_run_id == payload["node_run_id"],
|
||||
)
|
||||
)
|
||||
).one_or_none()
|
||||
if row is None:
|
||||
raise ValueError("node execution metadata not found")
|
||||
node_run, run, version, storage, workspace, schedule = row
|
||||
if storage.object_status != "available":
|
||||
raise ValueError("stable version artifact is not available")
|
||||
if storage.storage_backend != settings.storage_backend:
|
||||
raise ValueError(
|
||||
"稳定版本产物的存储后端与当前运行后端不一致"
|
||||
)
|
||||
if not storage.bucket_name or not storage.object_key:
|
||||
raise ValueError("stable version artifact location is incomplete")
|
||||
user_id = run.triggered_by or schedule.created_by
|
||||
# P0-5 / C1: re-verify the user is still active. ``create_scheduled_run``
|
||||
# checked membership when the run was queued, but the user may
|
||||
# have been disabled or soft-deleted in the meantime (admin
|
||||
# action, offboarding). Skip the check for the synthetic SYSTEM_CRON
|
||||
# user — that row is a fixed admin baseline and never goes inactive.
|
||||
if user_id != SYSTEM_CRON_USER_ID:
|
||||
await self._assert_user_active(session, user_id)
|
||||
context = {
|
||||
"node_status": node_run.node_status,
|
||||
"workspace_id": run.workspace_id,
|
||||
"workspace_code": workspace.workspace_code,
|
||||
"user_id": user_id,
|
||||
"bucket_name": storage.bucket_name,
|
||||
"object_key": storage.object_key,
|
||||
"content_hash": version.content_hash,
|
||||
}
|
||||
logger.debug(
|
||||
"execution context loaded: node_run={} bucket={} object_key={}",
|
||||
payload["node_run_id"][-12:],
|
||||
context["bucket_name"],
|
||||
context["object_key"][-32:],
|
||||
)
|
||||
return context
|
||||
|
||||
async def _fallback_execution_context(
|
||||
self,
|
||||
payload: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
async with self.session_factory() as session:
|
||||
row = (
|
||||
await session.execute(
|
||||
select(ScheduleRuns, Workspaces, Schedules)
|
||||
.join(
|
||||
Workspaces,
|
||||
Workspaces.workspace_id == ScheduleRuns.workspace_id,
|
||||
)
|
||||
.join(
|
||||
Schedules,
|
||||
Schedules.schedule_id == ScheduleRuns.schedule_id,
|
||||
)
|
||||
.where(ScheduleRuns.run_id == payload["run_id"])
|
||||
)
|
||||
).one_or_none()
|
||||
if row is None:
|
||||
raise ValueError("schedule run execution context not found")
|
||||
run, workspace, schedule = row
|
||||
return {
|
||||
"workspace_id": run.workspace_id,
|
||||
"workspace_code": workspace.workspace_code,
|
||||
"user_id": run.triggered_by or schedule.created_by,
|
||||
}
|
||||
|
||||
async def _download_artifact(
|
||||
self,
|
||||
*,
|
||||
bucket_name: str,
|
||||
object_key: str,
|
||||
content_hash: str,
|
||||
) -> bytes:
|
||||
# Honor the artifact's actual bucket (P0-3 fix): the artifact may
|
||||
# live in ``Workspaces.artifact_bucket`` rather than the global
|
||||
# version bucket the default ``object_store`` is bound to.
|
||||
store = self._store_for(bucket_name)
|
||||
content = await store.get(object_key)
|
||||
if hashlib.sha256(content).hexdigest() != content_hash:
|
||||
raise ValueError("stable version artifact hash mismatch")
|
||||
logger.debug(
|
||||
"artifact downloaded: bucket={} object_key={} bytes={}",
|
||||
bucket_name,
|
||||
object_key[-32:],
|
||||
len(content),
|
||||
)
|
||||
return content
|
||||
|
||||
async def _upload_execution_artifacts(
|
||||
self,
|
||||
*,
|
||||
payload: dict[str, Any],
|
||||
context: dict[str, Any],
|
||||
result: ExecutionResult,
|
||||
) -> tuple[str | None, str | None, str | None]:
|
||||
log_id: str | None = None
|
||||
result_id: str | None = None
|
||||
upload_error: str | None = None
|
||||
try:
|
||||
log_object = await self.storage_client.create_object(
|
||||
workspace_id=context["workspace_id"],
|
||||
user_id=context["user_id"],
|
||||
usage_type="run_log",
|
||||
file_name=f"{payload['node_run_id']}.log",
|
||||
content_type="text/plain; charset=utf-8",
|
||||
content=result.logs,
|
||||
idempotency_key=f"{payload['node_run_id']}:log",
|
||||
)
|
||||
log_id = log_object["storage_object_id"]
|
||||
result_object = await self.storage_client.create_object(
|
||||
workspace_id=context["workspace_id"],
|
||||
user_id=context["user_id"],
|
||||
usage_type="run_result",
|
||||
file_name=result.result_file_name,
|
||||
content_type=result.result_content_type,
|
||||
content=result.result,
|
||||
idempotency_key=f"{payload['node_run_id']}:result",
|
||||
)
|
||||
result_id = result_object["storage_object_id"]
|
||||
except Exception as exc:
|
||||
upload_error = f"result upload failed: {exc}"[:2000]
|
||||
logger.exception("failed to upload node execution artifacts")
|
||||
if log_id and result_id and not upload_error:
|
||||
logger.info(
|
||||
"artifacts uploaded: log_id={} result_id={}",
|
||||
log_id[-12:],
|
||||
result_id[-12:],
|
||||
)
|
||||
return log_id, result_id, upload_error
|
||||
|
||||
@staticmethod
|
||||
async def _start_inbox(
|
||||
session,
|
||||
*,
|
||||
consumer_name: str,
|
||||
event_id: str,
|
||||
message_id: str,
|
||||
) -> tuple[ConsumerInbox, bool]:
|
||||
item = await session.get(
|
||||
ConsumerInbox,
|
||||
(consumer_name, event_id),
|
||||
with_for_update=True,
|
||||
)
|
||||
if item is not None and item.process_status == "succeeded":
|
||||
return item, False
|
||||
if item is None:
|
||||
item = ConsumerInbox(
|
||||
consumer_name=consumer_name,
|
||||
event_id=event_id,
|
||||
process_status="processing",
|
||||
message_id=message_id,
|
||||
)
|
||||
session.add(item)
|
||||
else:
|
||||
item.process_status = "processing"
|
||||
item.message_id = message_id
|
||||
item.error_message = None
|
||||
return item, True
|
||||
|
||||
@staticmethod
|
||||
def _finish_inbox(item: ConsumerInbox) -> None:
|
||||
item.process_status = "succeeded"
|
||||
item.processed_at = utcnow()
|
||||
item.error_message = None
|
||||
|
||||
|
||||
__all__ = ["NodeExecutor"]
|
||||
@@ -1,103 +1,14 @@
|
||||
import argparse
|
||||
import json
|
||||
import sys
|
||||
import traceback
|
||||
from pathlib import Path
|
||||
"""Compatibility shim — the runner moved to ``schedule.execution.runners.notebook``.
|
||||
|
||||
import nbformat
|
||||
from loguru import logger
|
||||
from nbclient import NotebookClient
|
||||
|
||||
|
||||
def emit_outputs(notebook: object) -> None:
|
||||
for cell in notebook.cells: # type: ignore[attr-defined]
|
||||
if cell.get("cell_type") != "code":
|
||||
continue
|
||||
for output in cell.get("outputs", []):
|
||||
output_type = output.get("output_type")
|
||||
if output_type == "stream":
|
||||
text = output.get("text", "")
|
||||
print(
|
||||
"".join(text) if isinstance(text, list) else str(text),
|
||||
end="",
|
||||
flush=True,
|
||||
)
|
||||
elif output_type == "error":
|
||||
print(
|
||||
f"{output.get('ename', 'Error')}: {output.get('evalue', '')}",
|
||||
file=sys.stderr,
|
||||
flush=True,
|
||||
)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--input", required=True)
|
||||
parser.add_argument("--output", required=True)
|
||||
parser.add_argument("--timeout", required=True, type=int)
|
||||
parser.add_argument(
|
||||
"--python-version",
|
||||
choices=("3.8", "3.10", "3.12"),
|
||||
default="3.12",
|
||||
)
|
||||
parser.add_argument("--arguments-json", default="[]")
|
||||
args = parser.parse_args()
|
||||
|
||||
source = Path(args.input)
|
||||
output = Path(args.output)
|
||||
arguments = json.loads(args.arguments_json)
|
||||
if not isinstance(arguments, list) or not all(
|
||||
isinstance(item, str) for item in arguments
|
||||
):
|
||||
raise ValueError("arguments-json must contain an array of strings")
|
||||
|
||||
logger.info(
|
||||
"notebook runner start: input={} timeout={}s python={} args={}",
|
||||
source.name,
|
||||
args.timeout,
|
||||
args.python_version,
|
||||
len(arguments),
|
||||
)
|
||||
notebook = nbformat.read(source, as_version=4)
|
||||
if arguments:
|
||||
notebook.cells.insert(
|
||||
0,
|
||||
nbformat.v4.new_code_cell(
|
||||
"import sys\n"
|
||||
f"sys.argv = {json.dumps([source.name, *arguments], ensure_ascii=False)}",
|
||||
metadata={"tags": ["injected-parameters"]},
|
||||
),
|
||||
)
|
||||
exit_code = 0
|
||||
try:
|
||||
kernel_name = f"python{args.python_version.replace('.', '')}"
|
||||
client = NotebookClient(
|
||||
notebook,
|
||||
timeout=max(1, args.timeout),
|
||||
kernel_name=kernel_name,
|
||||
allow_errors=False,
|
||||
)
|
||||
logger.debug(
|
||||
"notebook client created: kernel={} timeout={}s",
|
||||
kernel_name,
|
||||
max(1, args.timeout),
|
||||
)
|
||||
# No explicit cwd — the kernel inherits the parent's cwd, which the
|
||||
# scheduler sets to the staged artifact directory. Keeping it here
|
||||
# avoids any "cwd must exist" requirement on the host.
|
||||
client.execute()
|
||||
logger.info("notebook client execute done: input={}", source.name)
|
||||
except Exception as exc:
|
||||
logger.exception("notebook execute failed: input={}", source.name)
|
||||
traceback.print_exc()
|
||||
exit_code = 124 if "timeout" in type(exc).__name__.lower() else 1
|
||||
finally:
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
nbformat.write(notebook, output)
|
||||
emit_outputs(notebook)
|
||||
logger.debug("notebook output written: {}", output)
|
||||
raise SystemExit(exit_code)
|
||||
The schedule worker's ``_execute_notebook`` launches the notebook subprocess
|
||||
as ``python -m schedule.notebook_runner``. That ``-m`` string is a stable
|
||||
contract, so this shim re-exports ``main`` from the real module rather than
|
||||
growing a second copy.
|
||||
"""
|
||||
|
||||
from schedule.execution.runners.notebook import main
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
__all__ = ["main"]
|
||||
@@ -38,7 +38,7 @@ from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
||||
|
||||
from schedule.scheduling.orchestrator import DispatchOrchestrator
|
||||
from schedule.scheduling.scheduler import CronScheduler
|
||||
from schedule.worker import NodeExecutor
|
||||
from schedule.execution.worker import NodeExecutor
|
||||
|
||||
|
||||
class SchedulerService:
|
||||
|
||||
@@ -43,7 +43,7 @@ from sqlalchemy import select
|
||||
|
||||
from schedule.domain.context import TERMINAL_NODE_STATES
|
||||
from schedule.domain.execution import ExecutionResult
|
||||
from schedule.execution import execute_artifact
|
||||
from schedule.execution.runners.notebook import execute_artifact
|
||||
|
||||
NODE_EXECUTE_EVENT = schedule_event_type("job.node.execute")
|
||||
NODE_FINISHED_EVENT = schedule_event_type("job.node.finished")
|
||||
|
||||
@@ -32,7 +32,7 @@ def _make_executor() -> tuple[SimpleNamespace, MagicMock, MagicMock]:
|
||||
construction time; it must be returned untouched by ``_store_for`` for
|
||||
the global version bucket.
|
||||
"""
|
||||
from schedule.worker import NodeExecutor
|
||||
from schedule.execution.worker import NodeExecutor
|
||||
|
||||
default_store = MagicMock(name="default_store")
|
||||
storage_client = MagicMock(name="storage_client")
|
||||
|
||||
Reference in New Issue
Block a user