Demonstrate how an absolute timeout would terminate a script that exceeds its allowed run time.
- Adjust the absolute timeout to 10 seconds to ensure that the long-running script is terminated.
- Use scripts that better demonstrate the timeout and CPU time limits.
Code implementation:
import subprocess
import os
import shutil
import tempfile
from concurrent.futures import ThreadPoolExecutor, as_completed
def execute_script_in_docker(script: str, timeout: int, cpu_limit: int):
# Create a temporary directory for the script
temp_dir = tempfile.mkdtemp()
script_path = os.path.join(temp_dir, 'user_script.py')
try:
# Save the script to a file
with open(script_path, 'w') as f:
f.write(script)
# Define the Docker run command with restrictions
docker_command = [
'docker', 'run', '--rm',
'--cpus', '0.5', # Limit to 0.5 CPU cores
'--memory', '256m', # Limit memory usage to 256MB
'--ulimit', f'cpu={cpu_limit}', # Limit CPU time to specified seconds
'--stop-timeout', '5', # Wait 5 seconds for graceful shutdown
'--network', 'none', # Disable network access
'--read-only', # Set root filesystem as read-only
'--cap-drop=ALL', # Drop all capabilities
'--user', '1000:1000', # Run as non-root user
'-v', f"{temp_dir}:/scripts:ro", # Mount the script directory as read-only
'python:3.9', 'python', '/scripts/user_script.py'
]
# Run the Docker container with a timeout
result = subprocess.run(docker_command, capture_output=True, text=True, timeout=timeout)
return result.stdout, result.stderr
except subprocess.TimeoutExpired:
return None, f"Script execution exceeded the {timeout} seconds timeout."
except Exception as e:
return None, str(e)
finally:
# Clean up the temporary directory
shutil.rmtree(temp_dir)
def run_scripts_in_pool(scripts, timeout, cpu_limit, max_workers):
results = []
with ThreadPoolExecutor(max_workers=max_workers) as executor:
future_to_script = {executor.submit(execute_script_in_docker, script, timeout, cpu_limit): script for script in scripts}
for future in as_completed(future_to_script):
script = future_to_script[future]
try:
stdout, stderr = future.result()
results.append((script, stdout, stderr))
except Exception as e:
results.append((script, None, str(e)))
return results
# Example usage
user_scripts = [
"""
import time
print("Starting script 1...")
time.sleep(20) # Simulate a task that should be terminated due to timeout
print("Script 1 complete.")
""",
"""
print("Starting script 2...")
result = sum([1, 2, 3, 4])
print(f"Script 2 result: {result}")
""",
]
# Run the scripts with a 10 seconds timeout and a 15 seconds CPU time limit, using a pool of 2 workers
timeout = 10
cpu_limit = 15
max_workers = 2
results = run_scripts_in_pool(user_scripts, timeout, cpu_limit, max_workers)
for script, stdout, stderr in results:
print("SCRIPT:", script)
print("STDOUT:", stdout)
print("STDERR:", stderr)
- Timeout Adjustment: The timeout is now set to 10 seconds to demonstrate the absolute timeout behavior.
- Script Adjustment: The first script contains a
time.sleep(20)
call, which will cause it to exceed the 10-second timeout and be terminated.
- Absolute Timeout: The
subprocess.run
function with thetimeout
parameter ensures that the script is terminated if it exceeds the specified wall-clock time. - CPU Time Limit: Docker’s
--ulimit cpu
option ensures that the script does not exceed the specified CPU time.
By running the above example, you should see that the first script gets terminated due to the absolute timeout, demonstrating the effectiveness of both timeout mechanisms.