feat: schedule node add python version

This commit is contained in:
tao.chen
2026-08-11 15:56:05 +08:00
parent e28bc39328
commit fe7f1a744e
9 changed files with 128 additions and 21 deletions
+21 -6
View File
@@ -39,12 +39,14 @@ async def _execute_notebook(
artifact_name: str,
arguments: list[str],
timeout_seconds: int,
python_version: str = "3.12",
) -> ExecutionResult:
output = artifact.with_name(f"executed-{artifact_name}")
logger.debug(
"notebook exec start: artifact={} timeout={}s args={}",
"notebook exec start: artifact={} timeout={}s version={} args={}",
artifact_name,
timeout_seconds,
python_version,
len(arguments),
)
process = await asyncio.create_subprocess_exec(
@@ -59,6 +61,8 @@ async def _execute_notebook(
str(max(1, timeout_seconds)),
"--arguments-json",
json.dumps(arguments, ensure_ascii=False),
"--python-version",
python_version,
cwd=str(artifact.parent),
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.STDOUT,
@@ -76,9 +80,10 @@ async def _execute_notebook(
error_code = "NODE_TIMEOUT"
error_message = f"notebook exceeded timeout of {timeout_seconds} seconds"
logger.warning(
"notebook exec timed out: artifact={} timeout={}s",
"notebook exec timed out: artifact={} timeout={}s version={}",
artifact_name,
timeout_seconds,
python_version,
)
else:
exit_code = process.returncode
@@ -114,9 +119,10 @@ async def _execute_notebook(
error_message=error_message,
)
logger.info(
"notebook exec done: artifact={} status={} exit_code={}",
"notebook exec done: artifact={} status={} version={} exit_code={}",
artifact_name,
result.status,
python_version,
result.exit_code,
)
return result
@@ -127,11 +133,15 @@ async def _execute_python(
*,
arguments: list[str],
timeout_seconds: int,
python_version: str = "3.12",
) -> ExecutionResult:
# NOTE: The worker process runs a single interpreter, so .py scripts
# continue to use sys.executable. Switching interpreters is future work.
logger.debug(
"python exec start: artifact={} timeout={}s args={}",
"python exec start: artifact={} timeout={}s version={} args={}",
artifact.name,
timeout_seconds,
python_version,
len(arguments),
)
process = await asyncio.create_subprocess_exec(
@@ -156,9 +166,10 @@ async def _execute_python(
ensure_ascii=False,
).encode("utf-8")
logger.warning(
"python exec timed out: artifact={} timeout={}s",
"python exec timed out: artifact={} timeout={}s version={}",
artifact.name,
timeout_seconds,
python_version,
)
return ExecutionResult(
status="timed_out",
@@ -190,9 +201,10 @@ async def _execute_python(
error_message=message,
)
logger.info(
"python exec done: artifact={} status={} exit_code={}",
"python exec done: artifact={} status={} version={} exit_code={}",
artifact.name,
result.status,
python_version,
result.exit_code,
)
return result
@@ -207,6 +219,7 @@ async def execute_artifact(
artifact_path: str,
arguments: list[str],
timeout_seconds: int,
python_version: str = "3.12",
) -> ExecutionResult:
logger.debug(
"execute_artifact: run={} node={} script_type={}",
@@ -230,12 +243,14 @@ async def execute_artifact(
artifact_name=artifact_name,
arguments=arguments,
timeout_seconds=timeout_seconds,
python_version=python_version,
)
if script_type == "python":
return await _execute_python(
artifact,
arguments=arguments,
timeout_seconds=timeout_seconds,
python_version=python_version,
)
logger.error("unsupported script_type: {}", script_type)
raise ValueError(f"unsupported script_type: {script_type}")
+13 -9
View File
@@ -1,14 +1,11 @@
from __future__ import annotations
import argparse
import json
import sys
import traceback
from pathlib import Path
from loguru import logger
import nbformat
from loguru import logger
from nbclient import NotebookClient
@@ -27,8 +24,7 @@ def emit_outputs(notebook: object) -> None:
)
elif output_type == "error":
print(
f"{output.get('ename', 'Error')}: "
f"{output.get('evalue', '')}",
f"{output.get('ename', 'Error')}: {output.get('evalue', '')}",
file=sys.stderr,
flush=True,
)
@@ -39,6 +35,11 @@ def main() -> None:
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()
@@ -51,9 +52,10 @@ def main() -> None:
raise ValueError("arguments-json must contain an array of strings")
logger.info(
"notebook runner start: input={} timeout={}s args={}",
"notebook runner start: input={} timeout={}s python={} args={}",
source.name,
args.timeout,
args.python_version,
len(arguments),
)
notebook = nbformat.read(source, as_version=4)
@@ -68,14 +70,16 @@ def main() -> None:
)
exit_code = 0
try:
kernel_name = f"python{args.python_version.replace('.', '')}"
client = NotebookClient(
notebook,
timeout=max(1, args.timeout),
kernel_name="python3",
kernel_name=kernel_name,
allow_errors=False,
)
logger.debug(
"notebook client created: kernel=python3 timeout={}s",
"notebook client created: kernel={} timeout={}s",
kernel_name,
max(1, args.timeout),
)
# No explicit cwd — the kernel inherits the parent's cwd, which the
+28
View File
@@ -28,6 +28,7 @@ from common.db import session_scope
from common.db.models import (
ConsumerInbox,
ScheduleNodeRuns,
ScheduleNodes,
ScheduleRuns,
Schedules,
StorageObjects,
@@ -86,6 +87,9 @@ class NodeExecutor:
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"],
@@ -94,6 +98,7 @@ class NodeExecutor:
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()
@@ -250,6 +255,29 @@ class NodeExecutor:
)
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 _execution_context(
self,
payload: dict[str, Any],