Add test to check device, pid and threadid are same for input_fn, output_fn and predict_fn

Author: sachanub

test/integration/local/test_model_fn_context.py renamed to test/integration/local/test_model_context.py

@@ -20,11 +20,16 @@
 import pytest
 import requests
 import torch
+from concurrent.futures import ThreadPoolExecutor
+import csv
 
 from integration import model_gpu_context_dir
 
 BASE_URL = "http://0.0.0.0:8080/"
 PING_URL = BASE_URL + "ping"
+INVOCATION_URL = BASE_URL + "models/model/invoke"
+GPU_COUNT = torch.cuda.device_count()
+GPU_IDS_EXPECTED = [i for i in range(GPU_COUNT)]
 
 
 @pytest.fixture(scope="module", autouse=True)
@@ -34,7 +39,7 @@ def container(image_uri):
             pytest.skip("Skipping because tests running on CPU instance")
 
         command = (
-            "docker run --gpus=all "
+            "docker run --gpus=all -p 8080:8080 "
             "--name sagemaker-pytorch-inference-toolkit-context-test "
             "-v {}:/opt/ml/model "
             "{} serve"
@@ -60,16 +65,53 @@ def container(image_uri):
 
 
 def test_context_all_device_ids():
-    gpu_count = torch.cuda.device_count()
+    time.sleep(5)
 
-    gpu_ids_expected = [i for i in range(gpu_count)]
     gpu_ids_actual = []
-
-    with open(os.path.join(model_gpu_context_dir, 'code', 'device_info.txt'), 'r') as f:
+    with open(os.path.join(model_gpu_context_dir, 'code', 'model_fn_device_info.txt'), 'r') as f:
         for line in f:
             gpu_ids_actual.append(int(line))
 
     gpu_ids_actual = list(set(gpu_ids_actual))
     gpu_ids_actual.sort()
 
-    assert gpu_ids_actual == gpu_ids_expected
+    assert gpu_ids_actual == GPU_IDS_EXPECTED
+
+
+def test_same_pid_threadid():
+    time.sleep(5)
+    headers = {"Content-Type": "application/json"}
+    with ThreadPoolExecutor(max_workers=GPU_COUNT) as executor:
+        for i in range(GPU_COUNT):
+            executor.submit(send_request, b'input', headers)
+
+    input_fn_device_info = read_csv("input_fn_device_info.csv")
+    output_fn_device_info = read_csv("output_fn_device_info.csv")
+    predict_fn_device_info = read_csv("predict_fn_device_info.csv")
+
+    assert len(input_fn_device_info) == len(output_fn_device_info) == len(predict_fn_device_info)
+
+    for input_fn_row, output_fn_row, predict_fn_row in zip(
+        input_fn_device_info, output_fn_device_info, predict_fn_device_info
+    ):
+
+        device_id_input_fn, pid_input_fn, threadid_input_fn = input_fn_row
+        device_id_output_fn, pid_output_fn, threadid_output_fn = output_fn_row
+        device_id_predict_fn, pid_predict_fn, threadid_predict_fn = predict_fn_row
+
+        assert device_id_input_fn == device_id_output_fn == device_id_predict_fn
+        assert pid_input_fn == pid_output_fn == pid_predict_fn
+        assert threadid_input_fn == threadid_output_fn == threadid_predict_fn
+
+
+def send_request(input_data, headers):
+    requests.post(INVOCATION_URL, data=input_data, headers=headers)
+
+
+def read_csv(filename):
+    data = []
+    with open(os.path.join(model_gpu_context_dir, 'code', filename), 'r') as csv_file:
+        csv_reader = csv.reader(csv_file)
+        for row in csv_reader:
+            data.append(row)
+    return data

test/resources/model_gpu_context/code/inference.py

@@ -14,14 +14,50 @@
 
 import os
 import torch
+import csv
+import threading
 
-def model_fn(model_dir, context):
-    script_dir = os.path.dirname(os.path.abspath(__file__))
-    file_path = os.path.join(script_dir, "device_info.txt")
-    
+SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
+
+def create_model_fn_device_list(filename, context):
+    file_path = os.path.join(SCRIPT_DIR, filename)
     device = torch.device("cuda:" + str(context.system_properties.get("gpu_id")))
     device_str = str(device)[-1]
     with open(file_path, "a") as file:
         file.write(device_str + "\n")
+    return
+
+def create_device_info_csv(filename, context):
+    pid = os.getpid()
+    threadid = threading.current_thread().ident
+
+    device = torch.device("cuda:" + str(context.system_properties.get("gpu_id")))
+    device_id = str(device)[-1]
+
+    data = [device_id, pid, threadid]
+    file_exists = os.path.isfile(filename)
+
+    with open(filename, "a", newline="") as csv_file:
+        writer = csv.writer(csv_file)
+        writer.writerow(data)
 
+    return
+
+def model_fn(model_dir, context):
+    create_model_fn_device_list("model_fn_device_info.txt", context)
     return 'model'
+
+
+def input_fn(data, content_type ,context):
+    create_device_info_csv("input_fn_device_info.csv", context)
+    return data
+
+
+def predict_fn(data, model, context):
+    create_device_info_csv("predict_fn_device_info.csv", context)
+    return b'output'
+
+
+def output_fn(prediction, accept, context):
+    create_device_info_csv("output_fn_device_info.csv", context)
+    return prediction