"""
Docker Runtime for Kailash Python SDK.
This module implements a Docker-based runtime for executing workflows where each
node runs in a separate Docker container. This enables scalable, isolated, and
reproducible workflow execution.
Key features:
- Container isolation for each node
- Resource constraints for execution
- Network communication between nodes
- Volume mounting for data exchange
- Orchestration of workflow execution
- Observability and monitoring
"""
import json
import logging
import os
import subprocess
import sys
import tempfile
from pathlib import Path
from typing import Any
from kailash.nodes.base import Node
from kailash.runtime._time_limits import _validate_limits, arm_time_limits
from kailash.runtime.cancellation import CancellationToken
# BaseRuntime doesn't exist - we'll implement task tracking methods directly
from kailash.sdk_exceptions import (
HardTimeLimitExceeded,
NodeConfigurationError,
NodeExecutionError,
SoftTimeLimitExceeded,
)
from kailash.tracking.manager import TaskManager
from kailash.workflow.graph import Workflow
logger = logging.getLogger(__name__)
class DockerNodeWrapper:
"""
Wrapper for running a Kailash node in a Docker container.
This class handles:
- Dockerfile generation
- Image building
- Container execution
- I/O mapping
- Result extraction
"""
def __init__(
self,
node: Node,
node_id: str,
base_image: str = "python:3.11-slim",
work_dir: Path | None = None,
sdk_path: Path | None = None,
):
"""
Initialize a Docker node wrapper.
Args:
node: The Kailash node to containerize.
node_id: The ID of the node in the workflow.
base_image: Base Docker image to use.
work_dir: Working directory for Docker files. If None, a temp dir is created.
sdk_path: Path to the Kailash SDK source. If None, tries to determine it.
"""
self.node = node
self.node_id = node_id
self.base_image = base_image
# Create or use work directory
if work_dir:
self.work_dir = Path(work_dir)
self.work_dir.mkdir(parents=True, exist_ok=True)
self._created_temp_dir = False
else:
self.work_dir = Path(tempfile.mkdtemp(prefix=f"kailash_docker_{node_id}_"))
self._created_temp_dir = True
# Find SDK path if not provided
if sdk_path:
self.sdk_path = Path(sdk_path)
else:
# Try to find SDK path from node's module
module_path = Path(self.node.__class__.__module__.replace(".", "/"))
if module_path.parts and module_path.parts[0] == "kailash":
# Find the SDK path by walking up to find the src/kailash directory
module_file = sys.modules[self.node.__class__.__module__].__file__
if module_file:
file_path = Path(module_file)
for parent in file_path.parents:
if (parent / "kailash").exists() and "site-packages" not in str(
parent
):
self.sdk_path = parent.parent
break
if not hasattr(self, "sdk_path"):
raise NodeConfigurationError(
"Could not determine SDK path. Please provide it explicitly."
)
# Container properties
self.image_name = f"kailash-node-{self.node_id.lower()}"
self.container_name = f"kailash-{self.node_id.lower()}"
self.container_id = None
# I/O directories
self.input_dir = self.work_dir / "input"
self.output_dir = self.work_dir / "output"
self.input_dir.mkdir(exist_ok=True)
self.output_dir.mkdir(exist_ok=True)
def prepare_dockerfile(self) -> Path:
"""
Generate a Dockerfile for the node.
Returns:
Path to the generated Dockerfile.
"""
dockerfile_path = self.work_dir / "Dockerfile"
entrypoint_path = self.work_dir / "entrypoint.py"
node_config_path = self.work_dir / "node.json"
# Save node configuration
with open(node_config_path, "w") as f:
# Create a serializable representation of the node
node_data = {
"class": self.node.__class__.__name__,
"module": self.node.__class__.__module__,
"node_id": self.node_id,
"name": getattr(self.node, "name", self.node.__class__.__name__),
}
# Add parameters if available
if hasattr(self.node, "get_parameters"):
node_data["parameters"] = {}
for name, param in self.node.get_parameters().items():
node_data["parameters"][name] = {
"name": param.name,
"type": str(param.type),
"required": param.required,
"description": param.description,
}
json.dump(node_data, f, indent=2)
# Create entrypoint script
with open(entrypoint_path, "w") as f:
f.write(
"""#!/usr/bin/env python3
import os
import sys
import json
import logging
import importlib
from pathlib import Path
# Configure logging
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger("kailash_node")
def main():
\"\"\"Run a Kailash node in a Docker container.\"\"\"
logger.info("Starting node execution")
# Load node configuration
node_config_path = Path("/app/node.json")
with open(node_config_path, 'r') as f:
node_data = json.load(f)
logger.info(f"Loaded configuration for {node_data['class']} node")
# Load runtime inputs if available
input_path = Path("/data/inputs/json/inputs.json")
runtime_inputs = {}
if input_path.exists():
logger.info(f"Loading inputs from {input_path}")
with open(input_path, 'r') as f:
runtime_inputs = json.load(f)
# Dynamically import the node class
logger.info(f"Importing {node_data['module']}.{node_data['class']}")
try:
module = importlib.import_module(node_data['module'])
node_class = getattr(module, node_data['class'])
except (ImportError, AttributeError) as e:
logger.error(f"Failed to import node class: {e}")
return 1
# Create node instance
logger.info(f"Creating node instance: {node_data['name']}")
try:
node = node_class(name=node_data['name'])
except Exception as e:
logger.error(f"Failed to create node instance: {e}")
return 1
# Execute node
logger.info(f"Executing node with inputs: {list(runtime_inputs.keys())}")
try:
result = node.execute(**runtime_inputs)
logger.info("Node execution completed successfully")
except Exception as e:
logger.error(f"Node execution failed: {e}")
# Save error information
with open("/data/outputs/json/error.json", 'w') as f:
json.dump({
"error": str(e),
"type": e.__class__.__name__
}, f, indent=2)
return 1
# Save results
logger.info("Saving execution results")
try:
result_path = Path("/data/outputs/json/result.json")
with open(result_path, 'w') as f:
# Handle non-serializable objects with basic conversion
try:
json.dump(result, f, indent=2)
except TypeError:
logger.warning("Result not directly JSON serializable, converting to string")
json.dump({"result": str(result)}, f, indent=2)
logger.info(f"Results saved to {result_path}")
except Exception as e:
logger.error(f"Failed to save results: {e}")
return 1
logger.info(f"Node {node_data['name']} execution completed")
return 0
if __name__ == "__main__":
sys.exit(main())
"""
)
# Make entrypoint executable
os.chmod(entrypoint_path, 0o755)
# Create Dockerfile
with open(dockerfile_path, "w") as f:
f.write(
f"""FROM {self.base_image}
# Set working directory
WORKDIR /app
# Install system dependencies if needed
RUN apt-get update && apt-get install -y --no-install-recommends \\
make build-essential \\
&& apt-get clean \\
&& rm -rf /var/lib/apt/lists/*
# Copy Kailash SDK
COPY sdk /app/sdk
# Install the SDK and its dependencies
RUN pip install --no-cache-dir -e /app/sdk
# Copy node configuration and entrypoint
COPY node.json /app/node.json
COPY entrypoint.py /app/entrypoint.py
# Create data directories
RUN mkdir -p /examples/data/input /examples/data/output
# Set entrypoint
ENTRYPOINT ["/app/entrypoint.py"]
"""
)
# Create SDK directory
sdk_dir = self.work_dir / "sdk"
sdk_dir.mkdir(exist_ok=True)
# Copy SDK files
self._copy_sdk_files(sdk_dir)
return dockerfile_path
def _copy_sdk_files(self, sdk_dir: Path):
"""
Copy SDK files to the Docker build context.
Args:
sdk_dir: Destination directory for SDK files.
"""
# Copy source directory
if (self.sdk_path / "src").exists():
import shutil
# Copy src directory
src_dir = self.sdk_path / "src"
shutil.copytree(src_dir, sdk_dir / "src", dirs_exist_ok=True)
# Copy setup files
for setup_file in ["setup.py", "pyproject.toml"]:
if (self.sdk_path / setup_file).exists():
shutil.copy(self.sdk_path / setup_file, sdk_dir / setup_file)
else:
raise NodeConfigurationError(
f"SDK source directory not found at {self.sdk_path}/src"
)
def build_image(self) -> str:
"""
Build the Docker image for the node.
Returns:
The name of the built Docker image.
"""
# Ensure Dockerfile exists
if not (self.work_dir / "Dockerfile").exists():
self.prepare_dockerfile()
logger.info(f"Building Docker image for node {self.node_id}")
try:
subprocess.run(
["docker", "build", "-t", self.image_name, "."],
cwd=self.work_dir,
check=True,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
)
logger.info(f"Successfully built image: {self.image_name}")
return self.image_name
except subprocess.CalledProcessError as e:
error_msg = f"Failed to build Docker image for node {self.node_id}: {e.stderr.decode()}"
logger.error(error_msg)
raise RuntimeError(error_msg)
def prepare_inputs(self, inputs: dict[str, Any]):
"""
Prepare inputs for node execution.
Args:
inputs: The inputs to pass to the node.
"""
input_file = self.input_dir / "inputs.json"
with open(input_file, "w") as f:
json.dump(inputs, f, indent=2)
def run_container(
self,
network: str | None = None,
env_vars: dict[str, str] | None = None,
resource_limits: dict[str, str] | None = None,
) -> bool:
"""
Run the node in a Docker container.
Args:
network: Docker network to use.
env_vars: Environment variables to pass to the container.
resource_limits: Resource limits (memory, CPU) for the container.
Returns:
True if container ran successfully.
"""
logger.info(f"Running node {self.node_id} in Docker container")
# Build command
cmd = ["docker", "run", "--rm"]
# Add container name
cmd.extend(["--name", self.container_name])
# Add network if specified
if network:
cmd.extend(["--network", network])
# Add environment variables
if env_vars:
for key, value in env_vars.items():
cmd.extend(["-e", f"{key}={value}"])
# Add resource limits
if resource_limits:
if "memory" in resource_limits:
cmd.extend(["--memory", resource_limits["memory"]])
if "cpu" in resource_limits:
cmd.extend(["--cpus", resource_limits["cpu"]])
# Add volume mounts for data
cmd.extend(
[
"-v",
f"{self.input_dir.absolute()}:/examples/data/input",
"-v",
f"{self.output_dir.absolute()}:/examples/data/output",
]
)
# Use the image
cmd.append(self.image_name)
try:
result = subprocess.run(
cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE
)
# Result could be used for logging output if needed
_ = result
logger.info(f"Container for node {self.node_id} ran successfully")
return True
except subprocess.CalledProcessError as e:
error_msg = f"Container for node {self.node_id} failed: {e.stderr.decode()}"
logger.error(error_msg)
# Check if there's an error file
error_file = self.output_dir / "error.json"
if error_file.exists():
with open(error_file) as f:
error_data = json.load(f)
error_msg = f"Node execution error: {error_data.get('error', 'Unknown error')}"
raise NodeExecutionError(error_msg)
def get_results(self) -> dict[str, Any]:
"""
Get the results of node execution.
Returns:
The node execution results.
"""
result_file = self.output_dir / "result.json"
if result_file.exists():
with open(result_file) as f:
return json.load(f)
error_file = self.output_dir / "error.json"
if error_file.exists():
with open(error_file) as f:
error_data = json.load(f)
raise NodeExecutionError(
f"Node {self.node_id} execution failed: {error_data.get('error', 'Unknown error')}"
)
return {"error": "No result or error file found"}
def cleanup(self):
"""Clean up resources created for this node."""
if self._created_temp_dir and self.work_dir.exists():
import shutil
shutil.rmtree(self.work_dir)
[docs]
class DockerRuntime:
"""
Docker-based runtime for executing workflows.
This runtime executes each node in a separate Docker container,
handling dependencies, data passing, and workflow orchestration.
"""
[docs]
def __init__(
self,
base_image: str = "python:3.11-slim",
network_name: str = "kailash-network",
work_dir: str | None = None,
sdk_path: str | None = None,
resource_limits: dict[str, str] | None = None,
task_manager: TaskManager | None = None,
):
"""
Initialize the Docker runtime.
Args:
base_image: Base Docker image to use for nodes.
network_name: Docker network name for container communication.
work_dir: Working directory for Docker files.
sdk_path: Path to the Kailash SDK source.
resource_limits: Default resource limits for containers.
task_manager: Task manager for tracking workflow execution.
"""
self.task_manager = task_manager
self.base_image = base_image
self.network_name = network_name
self.resource_limits = resource_limits or {}
# Working directory
if work_dir:
self.work_dir = Path(work_dir)
self.work_dir.mkdir(parents=True, exist_ok=True)
else:
self.work_dir = Path(tempfile.mkdtemp(prefix="kailash_docker_runtime_"))
# SDK path
if sdk_path:
self.sdk_path = Path(sdk_path)
else:
# Try to find the SDK path
import kailash
kailash_path = Path(kailash.__file__).parent
# Check if we're in a development environment
if "site-packages" not in str(kailash_path):
# Development environment - use parent of src
self.sdk_path = kailash_path.parent.parent
else:
# Installed package - use package directory
self.sdk_path = kailash_path.parent
# Create Docker network
self._create_network()
# Track node wrappers
self.node_wrappers = {}
def _create_task_run(self, workflow: Workflow) -> str | None:
"""Create a task run if task manager is available."""
if self.task_manager:
return self.task_manager.create_run(workflow.name)
return None
def _update_task_status(
self, run_id: str | None, node_id: str, status: str, output: Any = None
):
"""Update task status if task manager is available."""
if self.task_manager and run_id:
# For now, just update run status - task tracking needs more setup
if status == "failed":
error_msg = (
output.get("error", "Unknown error") if output else "Unknown error"
)
self.task_manager.update_run_status(run_id, "failed", error_msg)
def _complete_task_run(self, run_id: str | None, status: str, result: Any = None):
"""Complete task run if task manager is available."""
if self.task_manager and run_id:
if status == "completed":
self.task_manager.update_run_status(run_id, "completed")
else:
error_msg = (
result.get("error", "Unknown error") if result else "Unknown error"
)
self.task_manager.update_run_status(run_id, "failed", error_msg)
def _create_network(self):
"""Create a Docker network for container communication."""
try:
subprocess.run(
["docker", "network", "create", self.network_name],
check=True,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
)
logger.info(f"Created Docker network: {self.network_name}")
except subprocess.CalledProcessError as e:
# Ignore if network already exists
if "already exists" not in e.stderr.decode():
error_msg = f"Failed to create Docker network: {e.stderr.decode()}"
logger.error(error_msg)
raise RuntimeError(error_msg)
[docs]
def execute(
self,
workflow: Workflow,
inputs: dict[str, dict[str, Any]] | None = None,
node_resource_limits: dict[str, dict[str, str]] | None = None,
*,
soft_time_limit: float | None = None,
time_limit: float | None = None,
**kwargs: Any,
) -> tuple[dict[str, dict[str, Any]], str | None]:
"""
Execute a workflow using Docker containers.
Args:
workflow: The workflow to execute.
inputs: The inputs for each node.
node_resource_limits: Resource limits for specific nodes.
soft_time_limit: Optional advisory deadline in seconds (#912
Shard 1 slot; out-of-process timer enforcement lands later).
time_limit: Optional unconditional kill deadline in seconds.
**kwargs: Forward-compatibility kwargs for additive #912 Shard 1
contract.
Returns:
Tuple of (execution_results, run_id).
"""
# #912 Shard 1: validate typed time-limit kwargs at the entry point.
_validate_limits(soft_time_limit, time_limit)
# #912 Shard 6: arm threading.Timer-based deadlines around the
# per-node container-execution loop. Docker exec is out-of-
# process — the timers cannot inject into a running `docker run`
# subprocess — so the practical contract is: between nodes, we
# check the cancellation token + hard-deadline flag and raise.
# Long-running single-node Docker workflows do NOT get
# interrupted mid-container; that's an inherent constraint of
# the out-of-process model and is documented on the kwarg
# docstring above. Multi-node Docker workflows DO get the
# full deadline contract per-node-boundary.
_has_time_limit = soft_time_limit is not None or time_limit is not None
cancellable = None
_attempt_token: CancellationToken | None = None
if _has_time_limit:
_attempt_token = CancellationToken()
cancellable = arm_time_limits(
_attempt_token,
soft_time_limit=soft_time_limit,
time_limit=time_limit,
)
# Create task run
run_id = self._create_task_run(workflow)
# Default inputs
inputs = inputs or {}
node_resource_limits = node_resource_limits or {}
try:
try:
# Validate workflow
workflow.validate(runtime_parameters=inputs)
# Get execution order
execution_order = workflow.get_execution_order()
# Track results
results = {}
# Prepare all node wrappers and build images
logger.info("Preparing Docker containers for workflow execution")
for node_id, node in workflow.nodes.items():
self.node_wrappers[node_id] = DockerNodeWrapper(
node=node,
node_id=node_id,
base_image=self.base_image,
work_dir=self.work_dir / node_id,
sdk_path=self.sdk_path,
)
# Build image
self.node_wrappers[node_id].build_image()
# Execute nodes in order
logger.info(f"Executing workflow in order: {execution_order}")
for node_id in execution_order:
# Per-node-boundary check: did the hard timer fire
# while we were running the previous container?
if cancellable is not None:
if cancellable.hard_deadline_reached:
raise HardTimeLimitExceeded(
f"workflow exceeded hard time limit "
f"(time_limit={cancellable.time_limit}s + "
f"grace_seconds={cancellable.grace_seconds}s)"
)
if (
_attempt_token is not None
and _attempt_token.is_cancelled
and cancellable.soft_time_limit is not None
):
raise SoftTimeLimitExceeded(
f"workflow exceeded soft time limit "
f"(soft_time_limit={cancellable.soft_time_limit}s)"
)
logger.info(f"Executing node: {node_id}")
# Get node wrapper
wrapper = self.node_wrappers[node_id]
# Update task status
self._update_task_status(run_id, node_id, "running")
# Get node inputs
node_inputs = inputs.get(node_id, {}).copy()
# Add inputs from upstream nodes
_conn_map = getattr(workflow.connections, "get", lambda *a: {})(
node_id, {}
)
for upstream_id, mapping in (
_conn_map.items() if isinstance(_conn_map, dict) else []
):
if upstream_id in results:
for dest_param, src_param in mapping.items():
if src_param in results[upstream_id]:
node_inputs[dest_param] = results[upstream_id][
src_param
]
# Prepare inputs
wrapper.prepare_inputs(node_inputs)
# Get resource limits for this node
resource_limits = None
if node_id in node_resource_limits:
resource_limits = node_resource_limits[node_id]
elif self.resource_limits:
resource_limits = self.resource_limits
# Run the container
success = wrapper.run_container(
network=self.network_name, resource_limits=resource_limits
)
# Get results
if success:
results[node_id] = wrapper.get_results()
self._update_task_status(
run_id, node_id, "completed", results[node_id]
)
else:
self._update_task_status(
run_id, node_id, "failed", {"error": "Execution failed"}
)
raise NodeExecutionError(f"Node {node_id} execution failed")
# Final post-completion poll (Shard 2 invariant 5).
if cancellable is not None:
if cancellable.hard_deadline_reached:
raise HardTimeLimitExceeded(
f"workflow exceeded hard time limit "
f"(time_limit={cancellable.time_limit}s + "
f"grace_seconds={cancellable.grace_seconds}s)"
)
if (
_attempt_token is not None
and _attempt_token.is_cancelled
and cancellable.soft_time_limit is not None
):
raise SoftTimeLimitExceeded(
f"workflow exceeded soft time limit "
f"(soft_time_limit={cancellable.soft_time_limit}s)"
)
# Mark run as completed
self._complete_task_run(run_id, "completed")
return results, run_id
except (SoftTimeLimitExceeded, HardTimeLimitExceeded):
# Typed deadline exceptions MUST propagate untouched.
self._complete_task_run(
run_id, "failed", {"error": "Time limit exceeded"}
)
raise
except Exception as e:
# Handle errors
logger.error(f"Workflow execution failed: {e}")
self._complete_task_run(run_id, "failed", {"error": str(e)})
raise
finally:
if cancellable is not None:
cancellable.disarm()
[docs]
def cleanup(self):
"""Clean up Docker resources."""
# Clean up node wrappers
for wrapper in self.node_wrappers.values():
wrapper.cleanup()
# Remove Docker network
try:
subprocess.run(
["docker", "network", "rm", self.network_name],
check=False,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
)
except Exception as e:
logger.warning(f"Failed to remove Docker network: {e}")
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.cleanup()