fix: Propagate context to puller download tasks (#32)

#### Motivation

For some reason, the `LoadModel` context is not passed via the puller's `ProcessLoadModelRequest` function to pullman's `Pull` function. This means that storage downloads won't be cancelled when the corresponding request is cancelled or times out.

#### Modifications

Add a `Context` argument to the puller's `ProcessLoadModelRequest` function, use in the call to pull manager's `Pull` instead of `context.TODO()`.

#### Result

In-progress model downloads from shared storage downloads will be cancelled when the corresponding request is cancelled or times out.

Signed-off-by: Nick Hill <[email protected]>

Author: njhill

model-mesh-mlserver-adapter/server/server.go

@@ -121,7 +121,7 @@ func (s *MLServerAdapterServer) LoadModel(ctx context.Context, req *mmesh.LoadMo
 
 	if s.AdapterConfig.UseEmbeddedPuller {
 		var pullerErr error
-		req, pullerErr = s.Puller.ProcessLoadModelRequest(req)
+		req, pullerErr = s.Puller.ProcessLoadModelRequest(ctx, req)
 		if pullerErr != nil {
 			log.Error(pullerErr, "Failed to pull model from storage")
 			return nil, pullerErr

model-mesh-ovms-adapter/server/server.go

@@ -117,7 +117,7 @@ func (s *OvmsAdapterServer) LoadModel(ctx context.Context, req *mmesh.LoadModelR
 
 	if s.AdapterConfig.UseEmbeddedPuller {
 		var pullerErr error
-		req, pullerErr = s.Puller.ProcessLoadModelRequest(req)
+		req, pullerErr = s.Puller.ProcessLoadModelRequest(ctx, req)
 		if pullerErr != nil {
 			log.Error(pullerErr, "Failed to pull model from storage")
 			return nil, pullerErr

model-mesh-triton-adapter/server/server.go

@@ -106,7 +106,7 @@ func (s *TritonAdapterServer) LoadModel(ctx context.Context, req *mmesh.LoadMode
 
 	if s.AdapterConfig.UseEmbeddedPuller {
 		var pullerErr error
-		req, pullerErr = s.Puller.ProcessLoadModelRequest(req)
+		req, pullerErr = s.Puller.ProcessLoadModelRequest(ctx, req)
 		if pullerErr != nil {
 			log.Error(pullerErr, "Failed to pull model from storage")
 			return nil, pullerErr

model-serving-puller/puller/puller.go

@@ -91,7 +91,7 @@ func NewPullerFromConfig(log logr.Logger, config *PullerConfiguration) *Puller {
 // - rewrite ModelPath to a local filesystem path
 // - rewrite ModelKey["schema_path"] to a local filesystem path
 // - add the size of the model on disk to ModelKey["disk_size_bytes"]
-func (s *Puller) ProcessLoadModelRequest(req *mmesh.LoadModelRequest) (*mmesh.LoadModelRequest, error) {
+func (s *Puller) ProcessLoadModelRequest(ctx context.Context, req *mmesh.LoadModelRequest) (*mmesh.LoadModelRequest, error) {
 	var modelKey ModelKeyInfo
 	if parseErr := json.Unmarshal([]byte(req.ModelKey), &modelKey); parseErr != nil {
 		return nil, fmt.Errorf("Invalid modelKey in LoadModelRequest. Error processing JSON '%s': %w", req.ModelKey, parseErr)
@@ -177,7 +177,7 @@ func (s *Puller) ProcessLoadModelRequest(req *mmesh.LoadModelRequest) (*mmesh.Lo
 		Directory:        modelDir,
 		Targets:          targets,
 	}
-	pullerErr := s.PullManager.Pull(context.TODO(), pullCommand)
+	pullerErr := s.PullManager.Pull(ctx, pullCommand)
 	if pullerErr != nil {
 		return nil, status.Errorf(status.Code(pullerErr), "Failed to pull model from storage due to error: %s", pullerErr)
 	}

model-serving-puller/puller/puller_test.go

@@ -14,6 +14,7 @@
 package puller
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"io/ioutil"
@@ -176,7 +177,7 @@ func Test_ProcessLoadModelRequest_Success_SingleFileModel(t *testing.T) {
 
 	mockPuller.EXPECT().Pull(gomock.Any(), eqPullCommand(&expectedPullCommand)).Return(nil).Times(1)
 
-	returnRequest, err := p.ProcessLoadModelRequest(request)
+	returnRequest, err := p.ProcessLoadModelRequest(context.Background(), request)
 	assert.Nil(t, err)
 	assert.Equal(t, expectedRequestRewrite, returnRequest)
 }
@@ -216,7 +217,7 @@ func Test_ProcessLoadModelRequest_Success_MultiFileModel(t *testing.T) {
 
 	mockPuller.EXPECT().Pull(gomock.Any(), eqPullCommand(&expectedPullCommand)).Return(nil).Times(1)
 
-	returnRequest, err := p.ProcessLoadModelRequest(request)
+	returnRequest, err := p.ProcessLoadModelRequest(context.Background(), request)
 	assert.Nil(t, err)
 	assert.Equal(t, expectedRequestRewrite, returnRequest)
 }
@@ -262,7 +263,7 @@ func Test_ProcessLoadModelRequest_SuccessWithSchema(t *testing.T) {
 
 	mockPuller.EXPECT().Pull(gomock.Any(), eqPullCommand(&expectedPullCommand)).Return(nil).Times(1)
 
-	returnRequest, err := p.ProcessLoadModelRequest(request)
+	returnRequest, err := p.ProcessLoadModelRequest(context.Background(), request)
 	assert.Nil(t, err)
 	assert.Equal(t, expectedRequestRewrite, returnRequest)
 }
@@ -302,7 +303,7 @@ func Test_ProcessLoadModelRequest_SuccessWithBucket(t *testing.T) {
 
 	mockPuller.EXPECT().Pull(gomock.Any(), eqPullCommand(&expectedPullCommand)).Return(nil).Times(1)
 
-	returnRequest, err := p.ProcessLoadModelRequest(request)
+	returnRequest, err := p.ProcessLoadModelRequest(context.Background(), request)
 	assert.Nil(t, err)
 	assert.Equal(t, expectedRequestRewrite, returnRequest)
 }
@@ -342,7 +343,7 @@ func Test_ProcessLoadModelRequest_SuccessNoBucket(t *testing.T) {
 
 	mockPuller.EXPECT().Pull(gomock.Any(), eqPullCommand(&expectedPullCommand)).Return(nil).Times(1)
 
-	returnRequest, err := p.ProcessLoadModelRequest(request)
+	returnRequest, err := p.ProcessLoadModelRequest(context.Background(), request)
 	assert.Nil(t, err)
 	assert.Equal(t, expectedRequestRewrite, returnRequest)
 }
@@ -382,7 +383,7 @@ func Test_ProcessLoadModelRequest_SuccessNoBucketNoStorageParams(t *testing.T) {
 
 	mockPuller.EXPECT().Pull(gomock.Any(), eqPullCommand(&expectedPullCommand)).Return(nil).Times(1)
 
-	returnRequest, err := p.ProcessLoadModelRequest(request)
+	returnRequest, err := p.ProcessLoadModelRequest(context.Background(), request)
 	assert.Nil(t, err)
 	assert.Equal(t, expectedRequestRewrite, returnRequest)
 }
@@ -419,7 +420,7 @@ func Test_ProcessLoadModelRequest_SuccessStorageTypeOnly(t *testing.T) {
 
 	mockPuller.EXPECT().Pull(gomock.Any(), eqPullCommand(&expectedPullCommand)).Return(nil).Times(1)
 
-	returnRequest, err := p.ProcessLoadModelRequest(request)
+	returnRequest, err := p.ProcessLoadModelRequest(context.Background(), request)
 	assert.Nil(t, err)
 	assert.Equal(t, expectedRequestRewrite, returnRequest)
 }
@@ -466,7 +467,7 @@ func Test_ProcessLoadModelRequest_DefaultStorageKey(t *testing.T) {
 
 	mockPuller.EXPECT().Pull(gomock.Any(), eqPullCommand(&expectedPullCommand)).Return(nil).Times(1)
 
-	returnRequest, err := p.ProcessLoadModelRequest(request)
+	returnRequest, err := p.ProcessLoadModelRequest(context.Background(), request)
 	assert.Nil(t, err)
 	assert.Equal(t, expectedRequestRewrite, returnRequest)
 }
@@ -504,7 +505,7 @@ func Test_ProcessLoadModelRequest_DefaultStorageKeyTyped(t *testing.T) {
 
 	mockPuller.EXPECT().Pull(gomock.Any(), eqPullCommand(&expectedPullCommand)).Return(nil).Times(1)
 
-	returnRequest, err := p.ProcessLoadModelRequest(request)
+	returnRequest, err := p.ProcessLoadModelRequest(context.Background(), request)
 	assert.Nil(t, err)
 	assert.Equal(t, expectedRequestRewrite, returnRequest)
 }
@@ -544,7 +545,7 @@ func Test_ProcessLoadModelRequest_StorageParamsOverrides(t *testing.T) {
 
 	mockPuller.EXPECT().Pull(gomock.Any(), eqPullCommand(&expectedPullCommand)).Return(nil).Times(1)
 
-	returnRequest, err := p.ProcessLoadModelRequest(request)
+	returnRequest, err := p.ProcessLoadModelRequest(context.Background(), request)
 	assert.Nil(t, err)
 	assert.Equal(t, expectedRequestRewrite, returnRequest)
 }
@@ -559,7 +560,7 @@ func Test_ProcessLoadModelRequest_FailInvalidModelKey(t *testing.T) {
 
 	p, _ := newPullerWithMock(t)
 
-	returnRequest, err := p.ProcessLoadModelRequest(request)
+	returnRequest, err := p.ProcessLoadModelRequest(context.Background(), request)
 	assert.Contains(t, err.Error(), "Invalid modelKey in LoadModelRequest")
 	assert.Nil(t, returnRequest)
 }
@@ -574,7 +575,7 @@ func Test_ProcessLoadModelRequest_FailInvalidSchemaPath(t *testing.T) {
 
 	p, _ := newPullerWithMock(t)
 
-	returnRequest, err := p.ProcessLoadModelRequest(request)
+	returnRequest, err := p.ProcessLoadModelRequest(context.Background(), request)
 	assert.Nil(t, returnRequest)
 	assert.Error(t, err)
 	assert.Contains(t, err.Error(), "Invalid modelKey in LoadModelRequest")
@@ -591,7 +592,7 @@ func Test_ProcessLoadModelRequest_FailMissingStorageKeyAndType(t *testing.T) {
 
 	p, _ := newPullerWithMock(t)
 
-	returnRequest, err := p.ProcessLoadModelRequest(request)
+	returnRequest, err := p.ProcessLoadModelRequest(context.Background(), request)
 	assert.Nil(t, returnRequest)
 	assert.EqualError(t, err, expectedError)
 }

model-serving-puller/server/server.go

@@ -129,7 +129,7 @@ func (s *PullerServer) loadModel(ctx context.Context, req *mmesh.LoadModelReques
 
 	// Pull the model from storage
 	var pullerErr error
-	req, pullerErr = s.puller.ProcessLoadModelRequest(req)
+	req, pullerErr = s.puller.ProcessLoadModelRequest(ctx, req)
 	if pullerErr != nil {
 		log.Error(pullerErr, "Failed to pull model from storage")
 		return nil, pullerErr