feat: delete pipelines function

Author: riccardopinosio

hugot.go

@@ -22,7 +22,6 @@ type Session struct {
 }
 
 func newSession(runtime string, opts ...options.WithOption) (*Session, error) {
-
 	parsedOptions := options.Defaults()
 	parsedOptions.Runtime = runtime
 	// Collect options into a struct, so they can be applied in the correct order later
@@ -173,6 +172,8 @@ func InitializePipeline[T pipelineBackends.Pipeline](p T, pipelineConfig pipelin
 	default:
 		return pipeline, name, fmt.Errorf("not implemented")
 	}
+
+	model.Pipelines[name] = pipeline
 	return pipeline, name, nil
 }
 
@@ -209,6 +210,59 @@ func GetPipeline[T pipelineBackends.Pipeline](s *Session, name string) (T, error
 	}
 }
 
+func ClosePipeline[T pipelineBackends.Pipeline](s *Session, name string) error {
+	var pipeline T
+	switch any(pipeline).(type) {
+	case *pipelines.TokenClassificationPipeline:
+		p, ok := s.tokenClassificationPipelines[name]
+		if ok {
+			model := p.Model
+			delete(s.tokenClassificationPipelines, name)
+			delete(model.Pipelines, name)
+			if len(model.Pipelines) == 0 {
+				delete(s.models, model.Path)
+				return model.Destroy()
+			}
+		}
+	case *pipelines.TextClassificationPipeline:
+		p, ok := s.textClassificationPipelines[name]
+		if ok {
+			model := p.Model
+			delete(s.textClassificationPipelines, name)
+			delete(model.Pipelines, name)
+			if len(model.Pipelines) == 0 {
+				delete(s.models, model.Path)
+				return model.Destroy()
+			}
+		}
+	case *pipelines.FeatureExtractionPipeline:
+		p, ok := s.featureExtractionPipelines[name]
+		if ok {
+			model := p.Model
+			delete(s.featureExtractionPipelines, name)
+			delete(model.Pipelines, name)
+			if len(model.Pipelines) == 0 {
+				delete(s.models, model.Path)
+				return model.Destroy()
+			}
+		}
+	case *pipelines.ZeroShotClassificationPipeline:
+		p, ok := s.zeroShotClassificationPipelines[name]
+		if ok {
+			model := p.Model
+			delete(s.zeroShotClassificationPipelines, name)
+			delete(model.Pipelines, name)
+			if len(model.Pipelines) == 0 {
+				delete(s.models, model.Path)
+				return model.Destroy()
+			}
+		}
+	default:
+		return errors.New("pipeline type not supported")
+	}
+	return nil
+}
+
 type pipelineNotFoundError struct {
 	pipelineName string
 }

hugot_ort_test.go

@@ -241,6 +241,17 @@ func TestNoSameNamePipelineORT(t *testing.T) {
 	noSameNamePipeline(t, session)
 }
 
+func TestClosePipelineORT(t *testing.T) {
+	opts := []options.WithOption{options.WithOnnxLibraryPath(onnxRuntimeSharedLibrary)}
+	session, err := NewORTSession(opts...)
+	check(t, err)
+	defer func(session *Session) {
+		destroyErr := session.Destroy()
+		check(t, destroyErr)
+	}(session)
+	destroyPipelines(t, session)
+}
+
 // Thread safety
 
 func TestThreadSafetyORT(t *testing.T) {

hugot_test.go

@@ -827,6 +827,47 @@ func noSameNamePipeline(t *testing.T, session *Session) {
 	assert.Error(t, err3)
 }
 
+// destroy pipelines
+
+func destroyPipelines(t *testing.T, session *Session) {
+	t.Helper()
+
+	modelPath := "./models/KnightsAnalytics_distilbert-NER"
+	configSimple := TokenClassificationConfig{
+		ModelPath: modelPath,
+		Name:      "testClosePipeline",
+		Options: []TokenClassificationOption{
+			pipelines.WithSimpleAggregation(),
+			pipelines.WithIgnoreLabels([]string{"O"}),
+		},
+	}
+	_, err2 := NewPipeline(session, configSimple)
+	if err2 != nil {
+		t.FailNow()
+	}
+
+	if len(session.models) != 1 {
+		t.Fatal("Session should have 1 model")
+	}
+
+	for _, model := range session.models {
+		if _, ok := model.Pipelines["testClosePipeline"]; !ok {
+			t.Fatal("Pipeline alias was not added to the model")
+		}
+	}
+
+	if err := ClosePipeline[*pipelines.TokenClassificationPipeline](session, "testClosePipeline"); err != nil {
+		t.Fatal(err)
+	}
+
+	if len(session.models) != 0 {
+		t.Fatal("Session should have 0 models")
+	}
+	if len(session.tokenClassificationPipelines) != 0 {
+		t.Fatal("Session should have 0 token classification pipelines")
+	}
+}
+
 // Thread safety
 
 func threadSafety(t *testing.T, session *Session) {

hugot_xla_test.go

@@ -209,6 +209,16 @@ func TestNoSameNamePipelineXLA(t *testing.T) {
 	noSameNamePipeline(t, session)
 }
 
+func TestDestroyPipelineXLA(t *testing.T) {
+	session, err := NewXLASession()
+	check(t, err)
+	defer func(session *Session) {
+		destroyErr := session.Destroy()
+		check(t, destroyErr)
+	}(session)
+	destroyPipelines(t, session)
+}
+
 // Thread safety
 
 func TestThreadSafetyXLA(t *testing.T) {

pipelineBackends/model.go

@@ -16,6 +16,7 @@ type Model struct {
 	InputsMeta   []InputOutputInfo
 	OutputsMeta  []InputOutputInfo
 	Destroy      func() error
+	Pipelines    map[string]Pipeline
 }
 
 func ReshapeOutput(input *[]float32, meta InputOutputInfo, paddingMask [][]bool, sequenceLength int) OutputArray {
@@ -88,6 +89,7 @@ func LoadModel(path string, onnxFilename string, options *options.Options) (*Mod
 	model := &Model{
 		Path:         path,
 		OnnxFilename: onnxFilename,
+		Pipelines:    make(map[string]Pipeline),
 	}
 
 	err := LoadOnnxModelBytes(model)