Skip to content

Instantly share code, notes, and snippets.

@Darkflib
Created May 27, 2024 15:18
Show Gist options
  • Save Darkflib/4ff9c27394fa3f8a809544ae2fbb22b7 to your computer and use it in GitHub Desktop.
Save Darkflib/4ff9c27394fa3f8a809544ae2fbb22b7 to your computer and use it in GitHub Desktop.

Demonstrate how an absolute timeout would terminate a script that exceeds its allowed run time.

Example

  1. Adjust the absolute timeout to 10 seconds to ensure that the long-running script is terminated.
  2. 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)

Explanation of Code

  1. Timeout Adjustment: The timeout is now set to 10 seconds to demonstrate the absolute timeout behavior.
  2. Script Adjustment: The first script contains a time.sleep(20) call, which will cause it to exceed the 10-second timeout and be terminated.

Key Points

  • Absolute Timeout: The subprocess.run function with the timeout 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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment