Implement Prediction Job #751

HonakerM · 2024-07-25T11:57:00Z

What this PR does / why we need it:
This PR refactors the model trainer abstraction to be a more generic "job" abstraction. With this job abstraction we are able to implement both inferences and trainings in the background and poll their status/results.

This type of functionality has multiple use-cases inside Caikit with a few listed below:

Executing a long running inference which could be impacted by network issues
Tracking and managing inferences while they're running e.g. polling the status of an inference or cancelling a request.
Better support coordination between pods to support things like load distribution

This PR creates new http and grpc endpoints for the job predictions:

HTTP

POST /api/v1/task/<task>/job - create
GET /api/v1/task/<task>/job - get status
GET /api/v1/task/<task>/job/results - get results
DELETE /api/v1/task/<task>/job - cancel

GRPC:

caikit.runtime.SampleLib.SampleLibJobService.<task>StartPredictionJob - create job
caikit.runtime.SampleLib.SampleLibJobService.<task>PredictionJobStatus - get status
caikit.runtime.SampleLib.SampleLibJobService.<task>GetPredictionJobResult - get results
caikit.runtime.SampleLib.SampleLibJobService.<task>CancelPredictionJob - cancel

Special notes for your reviewer:

If applicable:

this PR contains documentation
this PR contains unit tests
this PR has been tested for backwards compatibility

Signed-off-by: Michael Honaker <Michael.Honaker@ibm.com>

caikit/core/data_model/training_status.py

caikit/interfaces/runtime/data_model/package.py

Signed-off-by: Michael Honaker <Michael.Honaker@ibm.com>

gabe-l-hart

A pile of NITs still (mostly copy-pasta) and a few structural things I still think need to be addressed to avoid API breaking changes:

Preserve kwarg names for ModelFuture to avoid breaking existing trainers
Don't switch 200 -> 202 (I kind of hate this one, but I think it's important)
Let's see if we can solve the py-to-proto duplicate import issue too!

caikit/core/data_model/training_status.py

caikit/core/data_model/__init__.py

caikit/core/model_management/local_job_predictor.py

caikit/core/model_manager.py

tests/runtime/http_server/test_http_server.py

tests/runtime/servicers/test_training_management_servicer.py

tests/runtime/test_grpc_server.py

Signed-off-by: Michael Honaker <Michael.Honaker@ibm.com>

tests/core/helpers.py

gabe-l-hart

Just a couple more NITs and a small outstanding list:

We need to figure out if changing TrainingStatus -> JobStatus in the TrainingStatusResponse is an API breaking change
job_id -> prediction_id in data model to mirror training_id for training jobs
Duplicate log code search (we should really put in a CI check for this)

caikit/core/data_model/prediction_status.py

caikit/core/model_management/job_predictor_base.py

gabe-l-hart · 2024-08-26T15:38:49Z

caikit/interfaces/runtime/data_model/training_management.py

@@ -43,7 +43,7 @@ class ModelPointer(DataObjectBase):
 @dataobject(RUNTIME_PACKAGE)
 class TrainingStatusResponse(DataObjectBase):
    training_id: Annotated[str, FieldNumber(1)]
-    state: Annotated[TrainingStatus, FieldNumber(2)]
+    state: Annotated[JobStatus, FieldNumber(2)]


Have we figured out whether this will cause an overall API breaking change by changing the name of the protobuf?

caikit/core/model_manager.py

gabe-l-hart · 2024-08-26T15:49:31Z

caikit/interfaces/runtime/data_model/prediction_job_management.py

+class PredictionJobInfoRequest(DataObjectBase):
+    """DataModel to request information about a PredictionJob"""
+
+    job_id: Annotated[str, FieldNumber(1)]


Since we're officially differentiating training_id and prediction_id elsewhere, I think maybe this should be named prediction_id instead of job_id? That would also mirror TrainingInfoRequest.

That's a good point. I'll make the update. I originally thought of making it more generic but since training is specific predictions should be as well

gabe-l-hart · 2024-08-26T15:53:21Z

caikit/runtime/servicers/global_predict_servicer.py

+            self.notify_backends_with_context(model_id, context)
+
+            # Retrieve the model from the model manager
+            log.debug("<RUN52259029D>", "Retrieving model '%s'", model_id)


Duplicate log code! Let's make sure we do a grep for these. You could also just remove them on any messages debug or lower

gabe-l-hart · 2024-08-26T15:55:02Z

caikit/runtime/servicers/prediction_job_management_servicer.py

+    ## Interface-agnostic entrypoints ##
+    ####################################
+
+    def get_job_result(self, job_id: str) -> DataObjectBase:


If we do job_id -> prediction_id in the interface, I think these should all change to <verb>_prediction_<noun>

tests/core/helpers.py

tests/core/model_management/test_local_job_predictor.py

caikit/runtime/http_server/http_server.py

Co-authored-by: Gabe Goodhart <gabe.l.hart@gmail.com> Signed-off-by: Michael Honaker <37811263+HonakerM@users.noreply.github.com>

Signed-off-by: Michael Honaker <Michael.Honaker@ibm.com>

gabe-l-hart

ONE. MORE. NIT!

gabe-l-hart · 2024-08-26T16:53:39Z

caikit/core/model_manager.py

+        Args:
+            future_id (str): The ID string from the original prediction
+                submission's ModelFuture
+
+        Returns:
+            model_future (JobPredictorFutureBase): The future handle
+                to the job which holds the status of the in-flight prediction.


Suggested change

Args:

future_id (str): The ID string from the original prediction

submission's ModelFuture

Returns:

model_future (JobPredictorFutureBase): The future handle

to the job which holds the status of the in-flight prediction.

Args:

prediction_id (str): The ID string from the original prediction

submission's ModelFuture

Returns:

prediction_future (JobPredictorFutureBase): The future handle

to the job which holds the status of the in-flight prediction.

Signed-off-by: Michael Honaker <Michael.Honaker@ibm.com>

gabe-l-hart

Ship it! I agree that we should call this an API breaking change even though it's very minimal. We should spell out exactly what would break in the release notes.

Signed-off-by: Michael Honaker <Michael.Honaker@ibm.com>

HonakerM added 5 commits July 21, 2024 15:45

Refactor Model Traininer to use Background Base

6692c2b

Signed-off-by: Michael Honaker <Michael.Honaker@ibm.com>

Update Tests and Runtime to use new ModelTrainer

8bdeefe

Signed-off-by: Michael Honaker <Michael.Honaker@ibm.com>

Add Initial Inference Implementation to HTTP

c574a8c

Signed-off-by: Michael Honaker <Michael.Honaker@ibm.com>

Clean up left over training lines

dc6967e

Signed-off-by: Michael Honaker <Michael.Honaker@ibm.com>

Rename BackgroundInference to Job

7a79ec8

Signed-off-by: Michael Honaker <Michael.Honaker@ibm.com>

HonakerM force-pushed the model_background_task branch from 6efcfae to 815c02a Compare July 25, 2024 11:57

GRPC Implementation Attempt

cb69d49

Signed-off-by: Michael Honaker <Michael.Honaker@ibm.com>

HonakerM force-pushed the model_background_task branch from 815c02a to c0419bb Compare July 31, 2024 19:11

Initial Tests and Formatting and Docs

8848015

Signed-off-by: Michael Honaker <Michael.Honaker@ibm.com>

HonakerM force-pushed the model_background_task branch 2 times, most recently from e19b10d to 875f66b Compare July 31, 2024 20:56

Add Result Call to ModelTrainerFuture

ba034eb

Signed-off-by: Michael Honaker <Michael.Honaker@ibm.com>

HonakerM force-pushed the model_background_task branch from 875f66b to ba034eb Compare July 31, 2024 21:19

Fix Tests and Remove Unused config

d1315f1

Signed-off-by: Michael Honaker <Michael.Honaker@ibm.com>

HonakerM force-pushed the model_background_task branch from 394a6a4 to 6badeab Compare August 1, 2024 16:01

Add Local Predictor tests

a3a7fd2

Signed-off-by: Michael Honaker <Michael.Honaker@ibm.com>

HonakerM force-pushed the model_background_task branch from 6badeab to 1a10065 Compare August 2, 2024 14:46

Add way to disable prediction job generation

38b3ced

Signed-off-by: Michael Honaker <Michael.Honaker@ibm.com>

HonakerM force-pushed the model_background_task branch from 1a10065 to 38b3ced Compare August 2, 2024 14:47

Add Exception for invalid service configuration

c38fdaa

Signed-off-by: Michael Honaker <Michael.Honaker@ibm.com>

HonakerM force-pushed the model_background_task branch from c049d97 to c38fdaa Compare August 2, 2024 15:52

Update Service Exports

854323e

Signed-off-by: Michael Honaker <Michael.Honaker@ibm.com>

HonakerM force-pushed the model_background_task branch from e403758 to 854323e Compare August 4, 2024 17:43

HonakerM added 3 commits August 4, 2024 16:06

Update Generation to Correctly create RPC Json

818546f

Signed-off-by: Michael Honaker <Michael.Honaker@ibm.com>

Add Recheck to HTTP Job Test

ac1730a

Signed-off-by: Michael Honaker <Michael.Honaker@ibm.com>

Rename Prediction Jobs to test GRPC

e8d7a80

Signed-off-by: Michael Honaker <Michael.Honaker@ibm.com>

HonakerM commented Aug 5, 2024

View reviewed changes

caikit/core/data_model/training_status.py Show resolved Hide resolved

HonakerM commented Aug 5, 2024

View reviewed changes

caikit/interfaces/runtime/data_model/package.py Outdated Show resolved Hide resolved

Remove unneeded regex

6b2c851

Signed-off-by: Michael Honaker <Michael.Honaker@ibm.com>

HonakerM marked this pull request as ready for review August 5, 2024 13:22

Update Async Resposne to send 202

b4be1ee

Signed-off-by: Michael Honaker <Michael.Honaker@ibm.com>

HonakerM force-pushed the model_background_task branch from b49d449 to b4be1ee Compare August 15, 2024 17:34

Update 200 to 202 in tests

4b3f876

Signed-off-by: Michael Honaker <Michael.Honaker@ibm.com>

gabe-l-hart requested changes Aug 19, 2024

View reviewed changes

HonakerM added 3 commits August 19, 2024 14:59

Rename status and update 202

31e8279

Signed-off-by: Michael Honaker <Michael.Honaker@ibm.com>

Fix Comments and KWargs

1e5ec52

Signed-off-by: Michael Honaker <Michael.Honaker@ibm.com>

Address Review Comments

1db2be0

Signed-off-by: Michael Honaker <Michael.Honaker@ibm.com>

HonakerM force-pushed the model_background_task branch from 3416581 to 1db2be0 Compare August 19, 2024 19:50

HonakerM added 2 commits August 19, 2024 16:09

Revert Package NAme Change

c883cf6

Signed-off-by: Michael Honaker <Michael.Honaker@ibm.com>

Add fix for Proto3

e30724e

Signed-off-by: Michael Honaker <Michael.Honaker@ibm.com>

HonakerM force-pushed the model_background_task branch from 21c487f to e30724e Compare August 20, 2024 16:42

Fix bad comment

e8a31f8

Signed-off-by: Michael Honaker <Michael.Honaker@ibm.com>

HonakerM force-pushed the model_background_task branch from f682417 to e8a31f8 Compare August 20, 2024 17:11

HonakerM requested a review from gabe-l-hart August 20, 2024 17:59

Fix Comments and remove training

309514a

Signed-off-by: Michael Honaker <Michael.Honaker@ibm.com>

HonakerM requested a review from tozhangkai August 23, 2024 14:20

HonakerM commented Aug 26, 2024

View reviewed changes

tests/core/helpers.py Show resolved Hide resolved

gabe-l-hart requested changes Aug 26, 2024

View reviewed changes

HonakerM and others added 5 commits August 26, 2024 12:06

Update caikit/core/data_model/prediction_status.py

f8eae82

Co-authored-by: Gabe Goodhart <gabe.l.hart@gmail.com> Signed-off-by: Michael Honaker <37811263+HonakerM@users.noreply.github.com>

Update caikit/core/model_management/job_predictor_base.py

9b8cb65

Co-authored-by: Gabe Goodhart <gabe.l.hart@gmail.com> Signed-off-by: Michael Honaker <37811263+HonakerM@users.noreply.github.com>

Update caikit/core/model_manager.py

4245077

Co-authored-by: Gabe Goodhart <gabe.l.hart@gmail.com> Signed-off-by: Michael Honaker <37811263+HonakerM@users.noreply.github.com>

Update tests/core/model_management/test_local_job_predictor.py

7bdfdbc

Co-authored-by: Gabe Goodhart <gabe.l.hart@gmail.com> Signed-off-by: Michael Honaker <37811263+HonakerM@users.noreply.github.com>

Rename job_id and remove dup logs

c03ef83

Signed-off-by: Michael Honaker <Michael.Honaker@ibm.com>

HonakerM force-pushed the model_background_task branch from 0e8065e to c03ef83 Compare August 26, 2024 16:37

HonakerM requested a review from gabe-l-hart August 26, 2024 16:50

gabe-l-hart requested changes Aug 26, 2024

View reviewed changes

One more NIT Fix

897df89

Signed-off-by: Michael Honaker <Michael.Honaker@ibm.com>

gabe-l-hart approved these changes Aug 26, 2024

View reviewed changes

Update tests for new prediction

8d64e8d

Signed-off-by: Michael Honaker <Michael.Honaker@ibm.com>

gabe-l-hart merged commit bd6e8a8 into caikit:main Aug 26, 2024
8 checks passed

Implement Prediction Job #751

Implement Prediction Job #751

Uh oh!

Conversation

HonakerM commented Jul 25, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

gabe-l-hart left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

gabe-l-hart left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

gabe-l-hart Aug 26, 2024

Choose a reason for hiding this comment

Uh oh!

Uh oh!

gabe-l-hart Aug 26, 2024

Choose a reason for hiding this comment

Uh oh!

HonakerM Aug 26, 2024

Choose a reason for hiding this comment

Uh oh!

gabe-l-hart Aug 26, 2024

Choose a reason for hiding this comment

Uh oh!

gabe-l-hart Aug 26, 2024

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

gabe-l-hart left a comment

Choose a reason for hiding this comment

Uh oh!

gabe-l-hart Aug 26, 2024

Choose a reason for hiding this comment

Uh oh!

gabe-l-hart left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

HonakerM commented Jul 25, 2024 •

edited

Loading