Repository: caikit/caikit
Pull request #614 (2023-12-14 19:05:26)

♻️ don't wait for training jobs to finish when using http


<!--  Thanks for sending a pull request!  Here are some tips for you:
1. Make sure to read the Contributing Guide before submitting your PR: https://github.com/caikit/caikit/blob/main/CONTRIBUTING.md
2. If this PR closes another issue, add 'closes #<issue number>' somewhere in the PR summary. GitHub will automatically close that issue when this PR gets merged. Alternatively, adding 'refs #<issue number>' will not close the issue, but help provide the reviewer more context.-->

**What this PR does / why we need it**:

fix https://github.com/caikit/caikit/issues/613

**Special notes for your reviewer**:

**If applicable**:
- [ ] this PR contains documentation
- [x] this PR contains unit tests
- [ ] this PR has been tested for backwards compatibility



Patch:
diff --git a/caikit/runtime/http_server/http_server.py b/caikit/runtime/http_server/http_server.py
index 4d723466..856cc214 100644
--- a/caikit/runtime/http_server/http_server.py
+++ b/caikit/runtime/http_server/http_server.py
@@ -423,7 +423,7 @@ class RuntimeHTTPServer(RuntimeServerBase):
                     module=rpc.clz,
                     training_output_dir=None,  # pass None so that GTS picks up the config one # TODO: double-check? # noqa: E501
                     # context=context,
-                    wait=True,
+                    wait=False,
                 )
                 result = await loop.run_in_executor(None, call)
                 if response_data_object.supports_file_operations:


Test patch:
diff --git a/tests/runtime/http_server/test_http_server.py b/tests/runtime/http_server/test_http_server.py
index 040c1ee7..0104027f 100644
--- a/tests/runtime/http_server/test_http_server.py
+++ b/tests/runtime/http_server/test_http_server.py
@@ -34,6 +34,7 @@ import tls_test_tools
 
 # Local
 from caikit.core import MODEL_MANAGER, DataObjectBase, dataobject
+from caikit.core.data_model import TrainingStatus
 from caikit.core.model_management.multi_model_finder import MultiModelFinder
 from caikit.runtime import http_server
 from caikit.runtime.http_server.http_server import StreamEventTypes
@@ -1083,6 +1084,40 @@ def test_http_and_grpc_server_share_threadpool(
     assert runtime_grpc_server.thread_pool is runtime_http_server.thread_pool
 
 
+def test_train_long_running_sample_task(client, runtime_http_server):
+    """Test that with a long running training job, the request returns before the training completes"""
+    model_name = "sample_task_train"
+    json_input = {
+        "model_name": model_name,
+        "parameters": {
+            "training_data": {"data_stream": {"data": [{"number": 1}]}},
+            "batch_size": 42,
+            "sleep_time": 5,  # mimic long train time
+        },
+    }
+    training_response = client.post(
+        f"/api/v1/SampleTaskSampleModuleTrain",
+        json=json_input,
+    )
+
+    # assert training response received before training completed
+    training_json_response = json.loads(
+        training_response.content.decode(training_response.default_encoding)
+    )
+    assert training_response.status_code == 200, training_json_response
+    assert (training_id := training_json_response["training_id"])
+    assert training_json_response["model_name"] == model_name
+
+    # assert that the training is still running
+    model_future = MODEL_MANAGER.get_model_future(training_id)
+    assert model_future.get_info().status == TrainingStatus.RUNNING
+
+    # Cancel the training
+    model_future.cancel()
+    assert model_future.get_info().status == TrainingStatus.CANCELED
+    assert model_future.get_info().status.is_terminal
+
+
 def test_uvicorn_server_config_valid():
     """Make sure that arbitrary uvicorn configs can be passed through from
     runtime.http.server_config