15.2.2 Measure Resources Before and After Running Your UDF

The example shows how to use psutil to monitor the CPU and memory resources consumed by your user-defined function (UDF):

# Load the libraries
import psutil
import time

# Define your UDF
def my_udf():
    from sklearn.datasets import load_iris
    from sklearn.model_selection import train_test_split
    from sklearn.ensemble import RandomForestClassifier
    from sklearn.metrics import accuracy_score
    # Load dataset
    data = load_iris()
    X = data.data
    y = data.target
    # Split the data into training and testing sets
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
    # Create a Random Forest classifier
    model = RandomForestClassifier(n_estimators=100, random_state=42)
    # Train the model
    model.fit(X_train, y_train)

# Get the current process
process = psutil.Process()
process

# Record resource usage before running the UDF
cpu_before = process.cpu_times()
memory_before = process.memory_info().rss
wall_start = time.time()

# Run your UDF
my_udf()

# Record resource usage after running the UDF
cpu_after = process.cpu_times()
memory_after = process.memory_info().rss
wall_elapsed = time.time() - wall_start

# Calculate resource usage
cpu_seconds = (cpu_after.user - cpu_before.user) + (cpu_after.system - cpu_before.system)
memory_usage = memory_after - memory_before
 

print(f"CPU time used by UDF: {cpu_seconds:.4f} seconds")
print(f"CPU usage by UDF: {(cpu_seconds / wall_elapsed) * 100:.1f}%")
print(f"Memory usage by UDF: {memory_usage / (1024 * 1024)} MB")

Example 15-1 Listing the Example

# Load the libraries
>>> import psutil
>>> import time

# Define your UDF
>>> def my_udf():
...     from sklearn.datasets import load_iris
...     from sklearn.model_selection import train_test_split
...     from sklearn.ensemble import RandomForestClassifier
...     from sklearn.metrics import accuracy_score
...     # Load dataset
...     data = load_iris()
...     X = data.data
...     y = data.target
...     # Split the data into training and testing sets
...     X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2\
, random_state=42)
...     # Create a Random Forest classifier
...     model = RandomForestClassifier(n_estimators=100, random_state=42)
...     # Train the model
...     model.fit(X_train, y_train)
... 
... # Get the current process
... process = psutil.Process()
... process
... 
psutil.Process(pid=392009, name='python3', status='running')
>>> # Record resource usage before running the UDF
... cpu_before = process.cpu_times()
... memory_before = process.memory_info().rss
... wall_start = time.time()
... 
... # Run your UDF
... my_udf()
... 
... # Record resource usage after running the UDF
... cpu_after = process.cpu_times()
... memory_after = process.memory_info().rss
... wall_elapsed = time.time() - wall_start
... 
... # Calculate resource usage
... cpu_seconds = (cpu_after.user - cpu_before.user) + (cpu_after.system - cpu_before.system)
... memory_usage = memory_after - memory_before
...  
... 
... print(f"CPU time used by UDF: {cpu_seconds:.4f} seconds")
... print(f"CPU usage by UDF: {(cpu_seconds / wall_elapsed) * 100:.1f}%")
... print(f"Memory usage by UDF: {memory_usage / (1024 * 1024)} MB")
... 

CPU time used by UDF: 1.6200 seconds
CPU usage by UDF: 59.5%
Memory usage by UDF: 133.46875 MB