From 82e83bf253d6fc29bbecd7528650e5e6d3793531 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 21 Oct 2026 02:20:12 +0000 Subject: [PATCH] test(addons): move the host operation tests to Ginkgo The hostops tests used the testing package calls that forbidigo rejects. They now run as Ginkgo specs under one suite bootstrap. Every check and every per-operation case is kept. Assisted-by: Claude Code:claude-sonnet-5-5 go Signed-off-by: Ettore Di Giacinto --- core/services/addons/hostops/broker_test.go | 194 +++--- .../addons/hostops/engine_datasets_test.go | 243 ++++---- .../addons/hostops/engine_jobs_test.go | 568 ++++++++---------- .../addons/hostops/engine_models_test.go | 349 +++++------ .../addons/hostops/engine_progress_test.go | 259 ++++---- core/services/addons/hostops/registry_test.go | 129 ++-- core/services/addons/hostops/suite_test.go | 14 + 7 files changed, 799 insertions(+), 957 deletions(-) create mode 100644 core/services/addons/hostops/suite_test.go diff --git a/core/services/addons/hostops/broker_test.go b/core/services/addons/hostops/broker_test.go index 84215608a..ecb16bddb 100644 --- a/core/services/addons/hostops/broker_test.go +++ b/core/services/addons/hostops/broker_test.go @@ -7,10 +7,11 @@ import ( "encoding/json" "net/http" "net/http/httptest" - "testing" "time" "github.com/mudler/LocalAI/pkg/addons/contract" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" ) func env(addon, op, scope string, exp time.Time) contract.Envelope { @@ -27,22 +28,18 @@ type fixture struct { seen *Call } -func newFixture(t *testing.T, has func(user, feature string) bool) *fixture { - t.Helper() +func newFixture(has func(user, feature string) bool) *fixture { f := &fixture{inflight: NewInFlight(), seen: &Call{}} f.b = New(f.inflight, has, contract.DefaultLimits()) err := f.b.Register(Spec{Op: "finetune.get", Scope: "read", Feature: "fine_tuning", Handler: func(_ context.Context, c Call) (json.RawMessage, error) { *f.seen = c return json.RawMessage(`{"ok":true}`), nil }}) - if err != nil { - t.Fatal(err) - } - if err := f.b.Register(Spec{Op: "finetune.start", Scope: "apply", Feature: "fine_tuning", Handler: func(context.Context, Call) (json.RawMessage, error) { + Expect(err).ToNot(HaveOccurred()) + err = f.b.Register(Spec{Op: "finetune.start", Scope: "apply", Feature: "fine_tuning", Handler: func(context.Context, Call) (json.RawMessage, error) { return json.RawMessage(`{"id":"j1"}`), nil - }}); err != nil { - t.Fatal(err) - } + }}) + Expect(err).ToNot(HaveOccurred()) f.h = f.b.Handler("fine-tuning", "secret", []string{"finetune."}) return f } @@ -60,103 +57,100 @@ func (f *fixture) call(credential, operationID, op, body string) (int, contract. func always(string, string) bool { return true } -func TestAWrongCredentialIsRefused(t *testing.T) { - f := newFixture(t, always) - if code, _ := f.call("nope", "op1", "finetune.get", `{}`); code != http.StatusUnauthorized { - t.Fatalf("got %d", code) +var _ = Describe("Host operation broker", func() { + live := func(f *fixture, scope string) { + end := f.inflight.Begin(env("fine-tuning", "op1", scope, time.Now().Add(time.Minute))) + DeferCleanup(end) } -} -func TestACallWithoutALiveEnvelopeIsRefused(t *testing.T) { - f := newFixture(t, always) - code, out := f.call("secret", "op1", "finetune.get", `{}`) - if code != http.StatusForbidden || out.Error == nil || out.Error.Code != "no_live_operation" { - t.Fatalf("%d %+v", code, out.Error) - } -} + It("refuses a wrong credential", func() { + f := newFixture(always) + code, _ := f.call("nope", "op1", "finetune.get", `{}`) + Expect(code).To(Equal(http.StatusUnauthorized)) + }) -func TestAnExpiredAuthorityIsRefused(t *testing.T) { - f := newFixture(t, always) - end := f.inflight.Begin(env("fine-tuning", "op1", "read", time.Now().Add(-time.Second))) - defer end() - code, out := f.call("secret", "op1", "finetune.get", `{}`) - if code != http.StatusForbidden || out.Error.Code != "no_live_operation" { - t.Fatalf("%d %+v", code, out.Error) - } -} + It("refuses a call without a live envelope", func() { + f := newFixture(always) + code, out := f.call("secret", "op1", "finetune.get", `{}`) + Expect(code).To(Equal(http.StatusForbidden)) + Expect(out.Error).ToNot(BeNil()) + Expect(out.Error.Code).To(Equal("no_live_operation")) + }) -func TestAReadScopeCannotRunAnApplyOperation(t *testing.T) { - f := newFixture(t, always) - end := f.inflight.Begin(env("fine-tuning", "op1", "read", time.Now().Add(time.Minute))) - defer end() - code, out := f.call("secret", "op1", "finetune.start", `{}`) - if code != http.StatusForbidden || out.Error.Code != "scope" { - t.Fatalf("%d %+v", code, out.Error) - } -} + It("refuses an expired authority", func() { + f := newFixture(always) + end := f.inflight.Begin(env("fine-tuning", "op1", "read", time.Now().Add(-time.Second))) + DeferCleanup(end) + code, out := f.call("secret", "op1", "finetune.get", `{}`) + Expect(code).To(Equal(http.StatusForbidden)) + Expect(out.Error).ToNot(BeNil()) + Expect(out.Error.Code).To(Equal("no_live_operation")) + }) -func TestTheFeatureOfTheActorIsChecked(t *testing.T) { - f := newFixture(t, func(user, feature string) bool { return feature != "fine_tuning" }) - end := f.inflight.Begin(env("fine-tuning", "op1", "read", time.Now().Add(time.Minute))) - defer end() - code, out := f.call("secret", "op1", "finetune.get", `{}`) - if code != http.StatusForbidden || out.Error.Code != "feature" { - t.Fatalf("%d %+v", code, out.Error) - } -} + It("does not let a read scope run an apply operation", func() { + f := newFixture(always) + live(f, "read") + code, out := f.call("secret", "op1", "finetune.start", `{}`) + Expect(code).To(Equal(http.StatusForbidden)) + Expect(out.Error).ToNot(BeNil()) + Expect(out.Error.Code).To(Equal("scope")) + }) -func TestTheOperationRunsAsTheActorOfTheEnvelope(t *testing.T) { - f := newFixture(t, always) - end := f.inflight.Begin(env("fine-tuning", "op1", "read", time.Now().Add(time.Minute))) - defer end() - // An argument that names a user is refused, whatever the schema says. - code, out := f.call("secret", "op1", "finetune.get", `{"userId":"admin"}`) - if code != http.StatusBadRequest || out.Error.Code != "bad_request" { - t.Fatalf("%d %+v", code, out.Error) - } - if code, _ := f.call("secret", "op1", "finetune.get", `{"jobId":"j1"}`); code != http.StatusOK { - t.Fatalf("got %d", code) - } - if f.seen.Actor.UserID != "u1" || f.seen.AddonID != "fine-tuning" { - t.Fatalf("the operation did not run as the envelope's actor: %+v", f.seen) - } -} + It("checks the feature of the actor", func() { + f := newFixture(func(user, feature string) bool { return feature != "fine_tuning" }) + live(f, "read") + code, out := f.call("secret", "op1", "finetune.get", `{}`) + Expect(code).To(Equal(http.StatusForbidden)) + Expect(out.Error).ToNot(BeNil()) + Expect(out.Error.Code).To(Equal("feature")) + }) -func TestAnOperationOutsideTheAllowedPrefixIsRefused(t *testing.T) { - f := newFixture(t, always) - _ = f.b.Register(Spec{Op: "quantize.get", Scope: "read", Feature: "quantization", Handler: func(context.Context, Call) (json.RawMessage, error) { return nil, nil }}) - end := f.inflight.Begin(env("fine-tuning", "op1", "read", time.Now().Add(time.Minute))) - defer end() - if code, out := f.call("secret", "op1", "quantize.get", `{}`); code != http.StatusForbidden || out.Error.Code != "unknown_op" { - t.Fatalf("%d %+v", code, out.Error) - } -} + It("runs the operation as the actor of the envelope", func() { + f := newFixture(always) + live(f, "read") + // An argument that names a user is refused, whatever the schema says. + code, out := f.call("secret", "op1", "finetune.get", `{"userId":"admin"}`) + Expect(code).To(Equal(http.StatusBadRequest)) + Expect(out.Error).ToNot(BeNil()) + Expect(out.Error.Code).To(Equal("bad_request")) + code, _ = f.call("secret", "op1", "finetune.get", `{"jobId":"j1"}`) + Expect(code).To(Equal(http.StatusOK)) + Expect(f.seen.Actor.UserID).To(Equal("u1"), "the operation did not run as the envelope's actor: %+v", f.seen) + Expect(f.seen.AddonID).To(Equal("fine-tuning")) + }) -func TestAnOversizedBodyIsRefused(t *testing.T) { - f := newFixture(t, always) - end := f.inflight.Begin(env("fine-tuning", "op1", "read", time.Now().Add(time.Minute))) - defer end() - big := `{"x":"` + string(bytes.Repeat([]byte("a"), int(contract.DefaultLimits().PayloadBytes)+1)) + `"}` - if code, out := f.call("secret", "op1", "finetune.get", big); code != http.StatusRequestEntityTooLarge || out.Error.Code != "too_large" { - t.Fatalf("%d %+v", code, out.Error) - } -} + It("refuses an operation outside the allowed prefix", func() { + f := newFixture(always) + _ = f.b.Register(Spec{Op: "quantize.get", Scope: "read", Feature: "quantization", Handler: func(context.Context, Call) (json.RawMessage, error) { return nil, nil }}) + live(f, "read") + code, out := f.call("secret", "op1", "quantize.get", `{}`) + Expect(code).To(Equal(http.StatusForbidden)) + Expect(out.Error).ToNot(BeNil()) + Expect(out.Error.Code).To(Equal("unknown_op")) + }) -func TestAnIdempotentOperationRunsOncePerOperationID(t *testing.T) { - f := newFixture(t, always) - runs := 0 - _ = f.b.Register(Spec{Op: "finetune.once", Scope: "apply", Feature: "fine_tuning", Idempotent: true, Handler: func(context.Context, Call) (json.RawMessage, error) { - runs++ - return json.RawMessage(`{"id":"j9"}`), nil - }}) - end := f.inflight.Begin(env("fine-tuning", "op1", "apply", time.Now().Add(time.Minute))) - defer end() - for i := 0; i < 2; i++ { - if code, _ := f.call("secret", "op1", "finetune.once", `{}`); code != http.StatusOK { - t.Fatalf("call %d: %d", i, code) + It("refuses an oversized body", func() { + f := newFixture(always) + live(f, "read") + big := `{"x":"` + string(bytes.Repeat([]byte("a"), int(contract.DefaultLimits().PayloadBytes)+1)) + `"}` + code, out := f.call("secret", "op1", "finetune.get", big) + Expect(code).To(Equal(http.StatusRequestEntityTooLarge)) + Expect(out.Error).ToNot(BeNil()) + Expect(out.Error.Code).To(Equal("too_large")) + }) + + It("runs an idempotent operation once per operation ID", func() { + f := newFixture(always) + runs := 0 + _ = f.b.Register(Spec{Op: "finetune.once", Scope: "apply", Feature: "fine_tuning", Idempotent: true, Handler: func(context.Context, Call) (json.RawMessage, error) { + runs++ + return json.RawMessage(`{"id":"j9"}`), nil + }}) + live(f, "apply") + for i := 0; i < 2; i++ { + code, _ := f.call("secret", "op1", "finetune.once", `{}`) + Expect(code).To(Equal(http.StatusOK), "call %d", i) } - } - if runs != 1 { - t.Fatalf("a repeated operationId must not run the operation twice, ran %d", runs) - } -} + Expect(runs).To(Equal(1), "a repeated operationId must not run the operation twice") + }) +}) diff --git a/core/services/addons/hostops/engine_datasets_test.go b/core/services/addons/hostops/engine_datasets_test.go index b8a9f470a..123ac1b61 100644 --- a/core/services/addons/hostops/engine_datasets_test.go +++ b/core/services/addons/hostops/engine_datasets_test.go @@ -3,12 +3,12 @@ package hostops import ( "encoding/json" - "errors" "fmt" "strings" - "testing" "github.com/mudler/LocalAI/pkg/addons/contract" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" ) type fakeDatasets struct { @@ -37,160 +37,123 @@ func (f *fakeDatasets) Resolve(ref string) (string, bool) { return p, ok } -func datasetArgs(t *testing.T, name, content string) string { - t.Helper() +func datasetArgs(name, content string) string { b, err := json.Marshal(map[string]string{"name": name, "content": content}) - if err != nil { - t.Fatal(err) - } + Expect(err).ToNot(HaveOccurred()) return string(b) } -func datasetBroker(t *testing.T, ds DatasetStore) (*Broker, Spec) { - t.Helper() +func datasetBroker(ds DatasetStore) (*Broker, Spec) { b := New(NewInFlight(), func(string, string) bool { return true }, contract.DefaultLimits()) - if err := RegisterEngineDatasets(b, ds); err != nil { - t.Fatal(err) - } + Expect(RegisterEngineDatasets(b, ds)).To(Succeed()) return b, b.specs["finetune.dataset"] } -func TestDatasetStoresTheContentAndReturnsARefNotAPath(t *testing.T) { - ds := &fakeDatasets{} - _, s := datasetBroker(t, ds) - raw, err := call(t, s, "named-user", "u1", datasetArgs(t, "train.jsonl", "{\"a\":1}\n")) - if err != nil { - t.Fatal(err) - } - var out struct { - DatasetRef string `json:"datasetRef"` - } - if err := json.Unmarshal(raw, &out); err != nil { - t.Fatal(err) - } - if out.DatasetRef == "" || strings.Contains(out.DatasetRef, "/") { - t.Fatalf("ref %q", out.DatasetRef) - } - if ds.gotName != "train.jsonl" || string(ds.gotData) != "{\"a\":1}\n" { - t.Fatalf("stored %q %q", ds.gotName, ds.gotData) - } - if strings.Contains(string(raw), "/srv") { - t.Fatalf("a path in the result: %s", raw) - } +type datasetResult struct { + DatasetRef string `json:"datasetRef"` } -func TestDatasetIsIdempotentAndFineTuningOnly(t *testing.T) { - _, s := datasetBroker(t, &fakeDatasets{}) - if !s.Idempotent || s.Scope != "apply" || s.Feature != "fine_tuning" { - t.Fatalf("%+v", s) - } -} +var _ = Describe("Dataset operation", func() { + It("stores the content and returns a ref, not a path", func() { + ds := &fakeDatasets{} + _, s := datasetBroker(ds) + raw, err := call(s, "named-user", "u1", datasetArgs("train.jsonl", "{\"a\":1}\n")) + Expect(err).ToNot(HaveOccurred()) + var out datasetResult + Expect(json.Unmarshal(raw, &out)).To(Succeed()) + Expect(out.DatasetRef).ToNot(BeEmpty()) + Expect(out.DatasetRef).ToNot(ContainSubstring("/")) + Expect(ds.gotName).To(Equal("train.jsonl")) + Expect(string(ds.gotData)).To(Equal("{\"a\":1}\n")) + Expect(string(raw)).ToNot(ContainSubstring("/srv"), "a path in the result") + }) -func TestDatasetRefusesAnEmptyOrOversizedContent(t *testing.T) { - ds := &fakeDatasets{} - _, s := datasetBroker(t, ds) - limit := 900 * 1024 - if _, err := call(t, s, "named-user", "u1", datasetArgs(t, "a.txt", strings.Repeat("x", limit))); err != nil { - t.Fatalf("exactly 900 KiB refused: %v", err) - } - for name, content := range map[string]string{ - "empty": "", - "one over": strings.Repeat("x", limit+1), - "far over": strings.Repeat("x", 2*limit), - } { - if _, err := call(t, s, "named-user", "u1", datasetArgs(t, "a.txt", content)); !errors.Is(err, ErrBadRequest) { - t.Fatalf("%s: want ErrBadRequest, got %v", name, err) - } - } - if ds.puts != 1 { - t.Fatalf("a refused dataset was stored: %d puts", ds.puts) - } -} + It("is idempotent and fine-tuning only", func() { + _, s := datasetBroker(&fakeDatasets{}) + Expect(s.Idempotent).To(BeTrue()) + Expect(s.Scope).To(Equal("apply")) + Expect(s.Feature).To(Equal("fine_tuning")) + }) -func TestDatasetRefusesANameWithADirectory(t *testing.T) { - ds := &fakeDatasets{} - _, s := datasetBroker(t, ds) - for _, name := range []string{"../x.jsonl", "a/b.jsonl", "/etc/x.jsonl", `a\b.jsonl`, ".."} { - if _, err := call(t, s, "named-user", "u1", datasetArgs(t, name, "[]")); !errors.Is(err, ErrBadRequest) { - t.Fatalf("%q: want ErrBadRequest, got %v", name, err) + It("refuses an empty or oversized content", func() { + ds := &fakeDatasets{} + _, s := datasetBroker(ds) + limit := 900 * 1024 + _, err := call(s, "named-user", "u1", datasetArgs("a.txt", strings.Repeat("x", limit))) + Expect(err).ToNot(HaveOccurred(), "exactly 900 KiB refused") + for name, content := range map[string]string{ + "empty": "", + "one over": strings.Repeat("x", limit+1), + "far over": strings.Repeat("x", 2*limit), + } { + _, err := call(s, "named-user", "u1", datasetArgs("a.txt", content)) + Expect(err).To(MatchError(ErrBadRequest), name) } - } - if ds.puts != 0 { - t.Fatalf("a refused name was stored") - } -} + Expect(ds.puts).To(Equal(1), "a refused dataset was stored") + }) -func TestDatasetRefusesAnExtensionThatIsNotATextDataset(t *testing.T) { - ds := &fakeDatasets{} - _, s := datasetBroker(t, ds) - for _, name := range []string{"x.exe", "x.sh", "x", "x.jsonl.exe", ".jsonl.bin"} { - if _, err := call(t, s, "named-user", "u1", datasetArgs(t, name, "[]")); !errors.Is(err, ErrBadRequest) { - t.Fatalf("%q: want ErrBadRequest, got %v", name, err) + It("refuses a name with a directory", func() { + ds := &fakeDatasets{} + _, s := datasetBroker(ds) + for _, name := range []string{"../x.jsonl", "a/b.jsonl", "/etc/x.jsonl", `a\b.jsonl`, ".."} { + _, err := call(s, "named-user", "u1", datasetArgs(name, "[]")) + Expect(err).To(MatchError(ErrBadRequest), name) } - } - for _, name := range []string{"x.jsonl", "x.json", "x.csv", "x.txt"} { - if _, err := call(t, s, "named-user", "u1", datasetArgs(t, name, "[]")); err != nil { - t.Fatalf("%q refused: %v", name, err) - } - } -} + Expect(ds.puts).To(Equal(0), "a refused name was stored") + }) -func TestDatasetRefusesContentThatIsNotUTF8(t *testing.T) { - ds := &fakeDatasets{} - _, s := datasetBroker(t, ds) - if _, err := call(t, s, "named-user", "u1", "{\"name\":\"x.csv\",\"content\":\"a,b\\n\xff\xfe,1\\n\"}"); !errors.Is(err, ErrBadRequest) { - t.Fatalf("got %v", err) - } - if ds.puts != 0 { - t.Fatalf("non-UTF-8 content was stored") - } -} + It("refuses an extension that is not a text dataset", func() { + ds := &fakeDatasets{} + _, s := datasetBroker(ds) + for _, name := range []string{"x.exe", "x.sh", "x", "x.jsonl.exe", ".jsonl.bin"} { + _, err := call(s, "named-user", "u1", datasetArgs(name, "[]")) + Expect(err).To(MatchError(ErrBadRequest), name) + } + for _, name := range []string{"x.jsonl", "x.json", "x.csv", "x.txt"} { + _, err := call(s, "named-user", "u1", datasetArgs(name, "[]")) + Expect(err).ToNot(HaveOccurred(), "%q refused", name) + } + }) -func TestResolveOnlyReturnsAReferenceTheHostIssuedToThatUser(t *testing.T) { - ds := &fakeDatasets{} - b, s := datasetBroker(t, ds) - raw, err := call(t, s, "named-user", "userA", datasetArgs(t, "x.jsonl", "[]")) - if err != nil { - t.Fatal(err) - } - var out struct { - DatasetRef string `json:"datasetRef"` - } - if err := json.Unmarshal(raw, &out); err != nil { - t.Fatal(err) - } - if _, ok := b.resolveDataset("userB", out.DatasetRef); ok { - t.Fatal("user B resolved the reference of user A") - } - if p, ok := b.resolveDataset("userA", out.DatasetRef); !ok || p == "" { - t.Fatalf("the issuing user cannot resolve its own reference: %q %v", p, ok) - } -} + It("refuses content that is not UTF-8", func() { + ds := &fakeDatasets{} + _, s := datasetBroker(ds) + _, err := call(s, "named-user", "u1", "{\"name\":\"x.csv\",\"content\":\"a,b\\n\xff\xfe,1\\n\"}") + Expect(err).To(MatchError(ErrBadRequest)) + Expect(ds.puts).To(Equal(0), "non-UTF-8 content was stored") + }) -func TestTheIssuedReferencesAreBoundedTo256(t *testing.T) { - ds := &fakeDatasets{} - b, s := datasetBroker(t, ds) - var first, last string - for i := range 257 { - raw, err := call(t, s, "named-user", "u1", datasetArgs(t, "x.jsonl", "[]")) - if err != nil { - t.Fatal(err) + It("resolves only a reference the host issued to that user", func() { + ds := &fakeDatasets{} + b, s := datasetBroker(ds) + raw, err := call(s, "named-user", "userA", datasetArgs("x.jsonl", "[]")) + Expect(err).ToNot(HaveOccurred()) + var out datasetResult + Expect(json.Unmarshal(raw, &out)).To(Succeed()) + _, ok := b.resolveDataset("userB", out.DatasetRef) + Expect(ok).To(BeFalse(), "user B resolved the reference of user A") + p, ok := b.resolveDataset("userA", out.DatasetRef) + Expect(ok).To(BeTrue(), "the issuing user cannot resolve its own reference") + Expect(p).ToNot(BeEmpty()) + }) + + It("bounds the issued references to 256", func() { + ds := &fakeDatasets{} + b, s := datasetBroker(ds) + var first, last string + for i := range 257 { + raw, err := call(s, "named-user", "u1", datasetArgs("x.jsonl", "[]")) + Expect(err).ToNot(HaveOccurred()) + var out datasetResult + Expect(json.Unmarshal(raw, &out)).To(Succeed()) + if i == 0 { + first = out.DatasetRef + } + last = out.DatasetRef } - var out struct { - DatasetRef string `json:"datasetRef"` - } - if err := json.Unmarshal(raw, &out); err != nil { - t.Fatal(err) - } - if i == 0 { - first = out.DatasetRef - } - last = out.DatasetRef - } - if _, ok := b.resolveDataset("u1", first); ok { - t.Fatal("the oldest reference is still issued after 257 datasets") - } - if _, ok := b.resolveDataset("u1", last); !ok { - t.Fatal("the newest reference is not issued") - } -} + _, ok := b.resolveDataset("u1", first) + Expect(ok).To(BeFalse(), "the oldest reference is still issued after 257 datasets") + _, ok = b.resolveDataset("u1", last) + Expect(ok).To(BeTrue(), "the newest reference is not issued") + }) +}) diff --git a/core/services/addons/hostops/engine_jobs_test.go b/core/services/addons/hostops/engine_jobs_test.go index 1710f467d..59ea1d199 100644 --- a/core/services/addons/hostops/engine_jobs_test.go +++ b/core/services/addons/hostops/engine_jobs_test.go @@ -5,15 +5,14 @@ import ( "bytes" "context" "encoding/json" - "errors" "net/http" "net/http/httptest" - "strings" - "testing" "time" "github.com/mudler/LocalAI/core/schema" "github.com/mudler/LocalAI/pkg/addons/contract" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" ) type fakeFT struct { @@ -78,185 +77,29 @@ type backends struct{} func (backends) FineTune() []string { return []string{"trl"} } func (backends) Quantization() []string { return []string{"llama-cpp-quantization"} } -func call(t *testing.T, spec Spec, mode, user, args string) (json.RawMessage, error) { - t.Helper() +func call(spec Spec, mode, user, args string) (json.RawMessage, error) { return spec.Handler(context.Background(), Call{AddonID: "fine-tuning", OperationID: "op1", Actor: contract.Actor{UserID: user, Mode: mode}, Arguments: json.RawMessage(args)}) } -func register(t *testing.T, ft FineTuner, q Quantizer) map[string]Spec { - t.Helper() +func register(ft FineTuner, q Quantizer) map[string]Spec { b := New(NewInFlight(), func(string, string) bool { return true }, contract.DefaultLimits()) - if err := RegisterEngineJobs(b, ft, q, backends{}); err != nil { - t.Fatal(err) - } + Expect(RegisterEngineJobs(b, ft, q, backends{})).To(Succeed()) return b.specs } -func specs(t *testing.T, ft FineTuner) map[string]Spec { - t.Helper() - return register(t, ft, nil) -} - -func TestStartPassesTheSameUserKeyAsTheRoute(t *testing.T) { - ft := &fakeFT{} - s := specs(t, ft) - if _, err := call(t, s["finetune.start"], "named-user", "u1", `{"model":"org/base","datasetSource":"org/data"}`); err != nil || ft.startUser != "u1" { - t.Fatalf("named user: %v %q", err, ft.startUser) - } - // Authentication off: the route passes "" so the job shows in the UI list. - if _, err := call(t, s["finetune.start"], "local-owner", "local-owner", `{"model":"org/base","datasetSource":"org/data"}`); err != nil || ft.startUser != "" { - t.Fatalf("local owner: %v %q", err, ft.startUser) - } -} - -func TestStartRefusesALocalPathAsModelOrDataset(t *testing.T) { - s := specs(t, &fakeFT{}) - for _, args := range []string{ - `{"model":"/etc/passwd","datasetSource":"org/data"}`, - `{"model":"../x","datasetSource":"org/data"}`, - `{"model":"org/base","datasetSource":"/var/data.jsonl"}`, - `{"model":"org\\base","datasetSource":"org/data"}`, - } { - if _, err := call(t, s["finetune.start"], "named-user", "u1", args); !errors.Is(err, ErrBadRequest) { - t.Fatalf("%s: want ErrBadRequest, got %v", args, err) - } - } -} - -func TestStartRefusesABackendThatIsNotAFineTuningBackend(t *testing.T) { - s := specs(t, &fakeFT{}) - if _, err := call(t, s["finetune.start"], "named-user", "u1", `{"model":"org/base","datasetSource":"org/data","backend":"llama-cpp"}`); !errors.Is(err, ErrBadRequest) { - t.Fatalf("got %v", err) - } -} - -func TestStartRefusesAFieldThatIsNotInTheToolSchema(t *testing.T) { - s := specs(t, &fakeFT{}) - if _, err := call(t, s["finetune.start"], "named-user", "u1", `{"model":"org/base","datasetSource":"org/data","extraOptions":{"x":"y"}}`); !errors.Is(err, ErrBadRequest) { - t.Fatalf("got %v", err) - } -} - -func TestGetOfAnotherUsersJobIsNotFound(t *testing.T) { - ft := &fakeFT{jobs: map[string]*schema.FineTuneJob{"j1": {ID: "j1", UserID: "other"}}} - s := specs(t, ft) - if _, err := call(t, s["finetune.get"], "named-user", "u1", `{"jobId":"j1"}`); !errors.Is(err, ErrNotFound) { - t.Fatalf("got %v", err) - } -} - -func TestListReturnsOnlyTheCallersJobs(t *testing.T) { - ft := &fakeFT{jobs: map[string]*schema.FineTuneJob{"a": {ID: "a", UserID: "u1"}, "b": {ID: "b", UserID: "u2"}}} - s := specs(t, ft) - out, err := call(t, s["finetune.list"], "named-user", "u1", `{}`) - if err != nil { - t.Fatal(err) - } - var got struct { - Jobs []struct { - ID string `json:"id"` - } `json:"jobs"` - } - _ = json.Unmarshal(out, &got) - if len(got.Jobs) != 1 || got.Jobs[0].ID != "a" { - t.Fatalf("a caller must see only its own jobs: %s", out) - } -} - -func TestAJobResultNeverCarriesAFilesystemPath(t *testing.T) { - ft := &fakeFT{jobs: map[string]*schema.FineTuneJob{"j1": {ID: "j1", UserID: "u1", OutputDir: "/data/fine-tune/j1", ExportStatus: "completed", ExportModelName: "my-model", Message: "saved to /data/fine-tune/j1/out"}}} - s := specs(t, ft) - out, _ := call(t, s["finetune.get"], "named-user", "u1", `{"jobId":"j1"}`) - if strings.Contains(string(out), "/data/") || strings.Contains(string(out), "/models/") { - t.Fatalf("a result must not leak a server path: %s", out) - } - if !strings.Contains(string(out), `"exportedModelName":"my-model"`) { - t.Fatalf("the registered model name is part of the result: %s", out) - } -} - -func TestQuantizeStartPassesTheSameUserKeyAsTheRoute(t *testing.T) { - q := &fakeQ{} - s := register(t, nil, q) - if _, err := call(t, s["quantize.start"], "named-user", "u1", `{"model":"org/base","quantizationType":"q8_0"}`); err != nil || q.startUser != "u1" { - t.Fatalf("named user: %v %q", err, q.startUser) - } - if _, err := call(t, s["quantize.start"], "local-owner", "local-owner", `{"model":"org/base"}`); err != nil || q.startUser != "" { - t.Fatalf("local owner: %v %q", err, q.startUser) - } - if got := q.started[0].QuantizationType; got != "q8_0" { - t.Fatalf("quantization type not passed: %q", got) - } -} - -func TestQuantizeStartRefusesALocalPathAsModel(t *testing.T) { - s := register(t, nil, &fakeQ{}) - for _, args := range []string{`{"model":"/models/x.gguf"}`, `{"model":"./x"}`, `{"model":"a\\b"}`} { - if _, err := call(t, s["quantize.start"], "named-user", "u1", args); !errors.Is(err, ErrBadRequest) { - t.Fatalf("%s: want ErrBadRequest, got %v", args, err) - } - } -} - -func TestQuantizeStartRefusesABackendThatIsNotAQuantizationBackend(t *testing.T) { - s := register(t, nil, &fakeQ{}) - if _, err := call(t, s["quantize.start"], "named-user", "u1", `{"model":"org/base","backend":"trl"}`); !errors.Is(err, ErrBadRequest) { - t.Fatalf("got %v", err) - } -} - -func TestQuantizeStartRefusesAFieldThatIsNotInTheToolSchema(t *testing.T) { - s := register(t, nil, &fakeQ{}) - if _, err := call(t, s["quantize.start"], "named-user", "u1", `{"model":"org/base","extraOptions":{"x":"y"}}`); !errors.Is(err, ErrBadRequest) { - t.Fatalf("got %v", err) - } -} - -func TestQuantizeGetOfAnotherUsersJobIsNotFound(t *testing.T) { - q := &fakeQ{jobs: map[string]*schema.QuantizationJob{"q1": {ID: "q1", UserID: "other", OutputFile: "/data/q1/m.gguf"}}} - s := register(t, nil, q) - if _, err := call(t, s["quantize.get"], "named-user", "u1", `{"jobId":"q1"}`); !errors.Is(err, ErrNotFound) { - t.Fatalf("got %v", err) - } -} - -func TestQuantizeListReturnsOnlyTheCallersJobsWithoutAPath(t *testing.T) { - q := &fakeQ{jobs: map[string]*schema.QuantizationJob{ - "a": {ID: "a", UserID: "u1", OutputFile: "/data/a/m.gguf", OutputDir: "/data/a"}, - "b": {ID: "b", UserID: "u2"}, - }} - s := register(t, nil, q) - out, err := call(t, s["quantize.list"], "named-user", "u1", `{}`) - if err != nil { - t.Fatal(err) - } - if strings.Contains(string(out), "/data/") || strings.Contains(string(out), `"id":"b"`) || !strings.Contains(string(out), `"id":"a"`) { - t.Fatalf("a caller sees only its own jobs and no path: %s", out) - } -} - -func TestNilServicesRegisterNoOperation(t *testing.T) { - b := New(NewInFlight(), always, contract.DefaultLimits()) - if err := RegisterEngineJobs(b, nil, nil, backends{}); err != nil { - t.Fatal(err) - } - if ops := b.Ops(); len(ops) != 0 { - t.Fatalf("no service, no operation: %v", ops) - } +func specs(ft FineTuner) map[string]Spec { + return register(ft, nil) } // serveEngine runs one engine operation through the broker's HTTP channel, as // the add-on process would call it, for a named user holding the features // that has allows. The jobs the fakes hold belong to u1 so that a call with // valid arguments answers 200 for an actor who holds the feature. -func serveEngine(t *testing.T, op, args string, has func(user, feature string) bool) (int, contract.HostResponse) { - t.Helper() +func serveEngine(op, args string, has func(user, feature string) bool) (int, contract.HostResponse) { ft := &fakeFT{jobs: map[string]*schema.FineTuneJob{"j1": {ID: "j1", UserID: "u1"}}} q := &fakeQ{jobs: map[string]*schema.QuantizationJob{"q1": {ID: "q1", UserID: "u1"}}} b := New(NewInFlight(), has, contract.DefaultLimits()) - if err := RegisterEngineJobs(b, ft, q, backends{}); err != nil { - t.Fatal(err) - } + Expect(RegisterEngineJobs(b, ft, q, backends{})).To(Succeed()) end := b.inflight.Begin(env("fine-tuning", "op1", "apply", time.Now().Add(time.Hour))) defer end() h := b.Handler("fine-tuning", "secret", []string{"finetune.", "quantize."}) @@ -270,154 +113,261 @@ func serveEngine(t *testing.T, op, args string, has func(user, feature string) b return rec.Code, out } -// TestEachEngineOperationNeedsItsFeature checks every engine operation against -// its own feature: refused without it, refused with only the other family's -// feature, and answered when the actor holds the operation's feature. -func TestEachEngineOperationNeedsItsFeature(t *testing.T) { - const fineTune, quant = "fine_tuning", "quantization" - cases := []struct { - op, args, feature, other string - }{ - {"finetune.start", `{"model":"org/base","datasetSource":"org/data"}`, fineTune, quant}, - {"finetune.get", `{"jobId":"j1"}`, fineTune, quant}, - {"finetune.list", `{}`, fineTune, quant}, - {"finetune.checkpoints", `{"jobId":"j1"}`, fineTune, quant}, - {"quantize.start", `{"model":"org/base"}`, quant, fineTune}, - {"quantize.get", `{"jobId":"q1"}`, quant, fineTune}, - {"quantize.list", `{}`, quant, fineTune}, - } - for _, c := range cases { - t.Run(c.op, func(t *testing.T) { - none := func(string, string) bool { return false } - if code, out := serveEngine(t, c.op, c.args, none); code != http.StatusForbidden || out.Error == nil || out.Error.Code != contract.HostFeature { - t.Fatalf("without any feature: %d %+v", code, out.Error) - } - otherOnly := func(_, f string) bool { return f == c.other } - if code, out := serveEngine(t, c.op, c.args, otherOnly); code != http.StatusForbidden || out.Error == nil || out.Error.Code != contract.HostFeature { - t.Fatalf("with only %s: %d %+v", c.other, code, out.Error) - } - ownOnly := func(_, f string) bool { return f == c.feature } - if code, out := serveEngine(t, c.op, c.args, ownOnly); code != http.StatusOK { - t.Fatalf("with %s: %d %+v", c.feature, code, out.Error) +var _ = Describe("Fine-tuning and quantization job operations", func() { + Describe("finetune", func() { + It("start passes the same user key as the route", func() { + ft := &fakeFT{} + s := specs(ft) + _, err := call(s["finetune.start"], "named-user", "u1", `{"model":"org/base","datasetSource":"org/data"}`) + Expect(err).ToNot(HaveOccurred()) + Expect(ft.startUser).To(Equal("u1"), "named user") + // Authentication off: the route passes "" so the job shows in the UI list. + _, err = call(s["finetune.start"], "local-owner", "local-owner", `{"model":"org/base","datasetSource":"org/data"}`) + Expect(err).ToNot(HaveOccurred()) + Expect(ft.startUser).To(BeEmpty(), "local owner") + }) + + It("start refuses a local path as model or dataset", func() { + s := specs(&fakeFT{}) + for _, args := range []string{ + `{"model":"/etc/passwd","datasetSource":"org/data"}`, + `{"model":"../x","datasetSource":"org/data"}`, + `{"model":"org/base","datasetSource":"/var/data.jsonl"}`, + `{"model":"org\\base","datasetSource":"org/data"}`, + } { + _, err := call(s["finetune.start"], "named-user", "u1", args) + Expect(err).To(MatchError(ErrBadRequest), args) } }) - } -} -// TestDatasetSourceMustBeAGalleryOrHuggingFaceID refuses a dataset reference -// that is not an id, names the field, and never reaches the service. -func TestDatasetSourceMustBeAGalleryOrHuggingFaceID(t *testing.T) { - bad := []string{ - "https://example.com/data.jsonl", // URL - "hf://org/data", // scheme - "file:///etc/passwd", // scheme and absolute path - "/var/data.jsonl", // absolute path - "../data", // parent segment - "org/../data", // parent segment inside an id - "org/..", // parent segment at the end - "org data", // space - "org/data/extra", // more than one separator - `C:\data`, // backslash - "", // empty - } - for _, src := range bad { - ft := &fakeFT{} - s := specs(t, ft) - args, _ := json.Marshal(map[string]string{"model": "org/base", "datasetSource": src}) - _, err := call(t, s["finetune.start"], "named-user", "u1", string(args)) - if !errors.Is(err, ErrBadRequest) || !strings.Contains(err.Error(), "datasetSource") { - t.Fatalf("%q: want a bad request naming datasetSource, got %v", src, err) - } - if len(ft.started) != 0 { - t.Fatalf("%q reached the service", src) - } - } -} + It("start refuses a backend that is not a fine-tuning backend", func() { + s := specs(&fakeFT{}) + _, err := call(s["finetune.start"], "named-user", "u1", `{"model":"org/base","datasetSource":"org/data","backend":"llama-cpp"}`) + Expect(err).To(MatchError(ErrBadRequest)) + }) -// TestDatasetSourceRefusalIsA400NamingTheField checks the answer the add-on -// process receives through the channel. -func TestDatasetSourceRefusalIsA400NamingTheField(t *testing.T) { - code, out := serveEngine(t, "finetune.start", `{"model":"org/base","datasetSource":"https://example.com/d.jsonl"}`, func(string, string) bool { return true }) - if code != http.StatusBadRequest || out.Error == nil || out.Error.Code != contract.HostBadRequest || !strings.Contains(out.Error.Message, "datasetSource") { - t.Fatalf("%d %+v", code, out.Error) - } -} + It("start refuses a field that is not in the tool schema", func() { + s := specs(&fakeFT{}) + _, err := call(s["finetune.start"], "named-user", "u1", `{"model":"org/base","datasetSource":"org/data","extraOptions":{"x":"y"}}`) + Expect(err).To(MatchError(ErrBadRequest)) + }) -// TestDatasetSourceAcceptsHuggingFaceAndGalleryIDs keeps the ids the UI uses. -func TestDatasetSourceAcceptsHuggingFaceAndGalleryIDs(t *testing.T) { - for _, src := range []string{"org/data", "HuggingFaceH4/ultrachat_200k", "my-dataset", "org/data-v1.0"} { - ft := &fakeFT{} - s := specs(t, ft) - args, _ := json.Marshal(map[string]string{"model": "org/base", "datasetSource": src}) - if _, err := call(t, s["finetune.start"], "named-user", "u1", string(args)); err != nil { - t.Fatalf("%q: %v", src, err) - } - } -} + It("get of another user's job is not found", func() { + ft := &fakeFT{jobs: map[string]*schema.FineTuneJob{"j1": {ID: "j1", UserID: "other"}}} + s := specs(ft) + _, err := call(s["finetune.get"], "named-user", "u1", `{"jobId":"j1"}`) + Expect(err).To(MatchError(ErrNotFound)) + }) -func TestStartAcceptsADatasetRefAsDatasetSource(t *testing.T) { - ft := &fakeFT{} - ds := &fakeDatasets{} - b := New(NewInFlight(), func(string, string) bool { return true }, contract.DefaultLimits()) - if err := RegisterEngineJobs(b, ft, nil, backends{}); err != nil { - t.Fatal(err) - } - if err := RegisterEngineDatasets(b, ds); err != nil { - t.Fatal(err) - } - raw, err := call(t, b.specs["finetune.dataset"], "named-user", "u1", datasetArgs(t, "x.jsonl", "[]")) - if err != nil { - t.Fatal(err) - } - var out struct { - DatasetRef string `json:"datasetRef"` - } - if err := json.Unmarshal(raw, &out); err != nil { - t.Fatal(err) - } - start := b.specs["finetune.start"] - if _, err := call(t, start, "named-user", "u1", `{"model":"org/base","datasetSource":"`+out.DatasetRef+`"}`); err != nil { - t.Fatal(err) - } - if got := ft.started[len(ft.started)-1].DatasetSource; got != ds.paths[out.DatasetRef] { - t.Fatalf("the service got %q, want the stored path", got) - } - if _, err := call(t, start, "named-user", "u2", `{"model":"org/base","datasetSource":"`+out.DatasetRef+`"}`); !errors.Is(err, ErrBadRequest) { - t.Fatalf("another user's reference: want ErrBadRequest, got %v", err) - } - if _, err := call(t, start, "named-user", "u1", `{"model":"org/base","datasetSource":"ffffffff-nope.jsonl"}`); !errors.Is(err, ErrBadRequest) { - t.Fatalf("an unknown reference: want ErrBadRequest, got %v", err) - } -} + It("list returns only the caller's jobs", func() { + ft := &fakeFT{jobs: map[string]*schema.FineTuneJob{"a": {ID: "a", UserID: "u1"}, "b": {ID: "b", UserID: "u2"}}} + s := specs(ft) + out, err := call(s["finetune.list"], "named-user", "u1", `{}`) + Expect(err).ToNot(HaveOccurred()) + var got struct { + Jobs []struct { + ID string `json:"id"` + } `json:"jobs"` + } + _ = json.Unmarshal(out, &got) + Expect(got.Jobs).To(HaveLen(1), "a caller must see only its own jobs: %s", out) + Expect(got.Jobs[0].ID).To(Equal("a")) + }) -func TestASecondUsersListDoesNotShowTheFirstUsersJobs(t *testing.T) { - ft := &fakeFT{jobs: map[string]*schema.FineTuneJob{"a": {ID: "a", UserID: "u1"}, "b": {ID: "b", UserID: "u2"}, "c": {ID: "c", UserID: "u1"}}} - s := specs(t, ft) - ids := func(user string) []string { - out, err := call(t, s["finetune.list"], "named-user", user, `{}`) - if err != nil { - t.Fatal(err) + It("a job result never carries a filesystem path", func() { + ft := &fakeFT{jobs: map[string]*schema.FineTuneJob{"j1": {ID: "j1", UserID: "u1", OutputDir: "/data/fine-tune/j1", ExportStatus: "completed", ExportModelName: "my-model", Message: "saved to /data/fine-tune/j1/out"}}} + s := specs(ft) + out, _ := call(s["finetune.get"], "named-user", "u1", `{"jobId":"j1"}`) + Expect(string(out)).ToNot(ContainSubstring("/data/"), "a result must not leak a server path") + Expect(string(out)).ToNot(ContainSubstring("/models/"), "a result must not leak a server path") + Expect(string(out)).To(ContainSubstring(`"exportedModelName":"my-model"`), "the registered model name is part of the result") + }) + + It("a second user's list does not show the first user's jobs", func() { + ft := &fakeFT{jobs: map[string]*schema.FineTuneJob{"a": {ID: "a", UserID: "u1"}, "b": {ID: "b", UserID: "u2"}, "c": {ID: "c", UserID: "u1"}}} + s := specs(ft) + ids := func(user string) []string { + out, err := call(s["finetune.list"], "named-user", user, `{}`) + Expect(err).ToNot(HaveOccurred()) + var got struct { + Jobs []struct { + ID string `json:"id"` + } `json:"jobs"` + } + _ = json.Unmarshal(out, &got) + var r []string + for _, j := range got.Jobs { + r = append(r, j.ID) + } + return r + } + Expect(ids("u2")).To(Equal([]string{"b"}), "user 2 must see only its own job") + Expect(ids("u1")).To(HaveLen(2), "user 1 must see its two jobs") + // A user with no job sees none, not the others'. + Expect(ids("u3")).To(BeEmpty(), "user 3 must see nothing") + }) + }) + + Describe("quantize", func() { + It("start passes the same user key as the route", func() { + q := &fakeQ{} + s := register(nil, q) + _, err := call(s["quantize.start"], "named-user", "u1", `{"model":"org/base","quantizationType":"q8_0"}`) + Expect(err).ToNot(HaveOccurred()) + Expect(q.startUser).To(Equal("u1"), "named user") + _, err = call(s["quantize.start"], "local-owner", "local-owner", `{"model":"org/base"}`) + Expect(err).ToNot(HaveOccurred()) + Expect(q.startUser).To(BeEmpty(), "local owner") + Expect(q.started[0].QuantizationType).To(Equal("q8_0"), "quantization type not passed") + }) + + It("start refuses a local path as model", func() { + s := register(nil, &fakeQ{}) + for _, args := range []string{`{"model":"/models/x.gguf"}`, `{"model":"./x"}`, `{"model":"a\\b"}`} { + _, err := call(s["quantize.start"], "named-user", "u1", args) + Expect(err).To(MatchError(ErrBadRequest), args) + } + }) + + It("start refuses a backend that is not a quantization backend", func() { + s := register(nil, &fakeQ{}) + _, err := call(s["quantize.start"], "named-user", "u1", `{"model":"org/base","backend":"trl"}`) + Expect(err).To(MatchError(ErrBadRequest)) + }) + + It("start refuses a field that is not in the tool schema", func() { + s := register(nil, &fakeQ{}) + _, err := call(s["quantize.start"], "named-user", "u1", `{"model":"org/base","extraOptions":{"x":"y"}}`) + Expect(err).To(MatchError(ErrBadRequest)) + }) + + It("get of another user's job is not found", func() { + q := &fakeQ{jobs: map[string]*schema.QuantizationJob{"q1": {ID: "q1", UserID: "other", OutputFile: "/data/q1/m.gguf"}}} + s := register(nil, q) + _, err := call(s["quantize.get"], "named-user", "u1", `{"jobId":"q1"}`) + Expect(err).To(MatchError(ErrNotFound)) + }) + + It("list returns only the caller's jobs without a path", func() { + q := &fakeQ{jobs: map[string]*schema.QuantizationJob{ + "a": {ID: "a", UserID: "u1", OutputFile: "/data/a/m.gguf", OutputDir: "/data/a"}, + "b": {ID: "b", UserID: "u2"}, + }} + s := register(nil, q) + out, err := call(s["quantize.list"], "named-user", "u1", `{}`) + Expect(err).ToNot(HaveOccurred()) + Expect(string(out)).ToNot(ContainSubstring("/data/"), "a caller sees no path") + Expect(string(out)).ToNot(ContainSubstring(`"id":"b"`), "a caller sees only its own jobs") + Expect(string(out)).To(ContainSubstring(`"id":"a"`), "a caller sees its own jobs") + }) + }) + + It("registers no operation for nil services", func() { + b := New(NewInFlight(), always, contract.DefaultLimits()) + Expect(RegisterEngineJobs(b, nil, nil, backends{})).To(Succeed()) + Expect(b.Ops()).To(BeEmpty(), "no service, no operation") + }) + + // Every engine operation is checked against its own feature: refused + // without it, refused with only the other family's feature, and answered + // when the actor holds the operation's feature. + Describe("feature gate", func() { + const fineTune, quant = "fine_tuning", "quantization" + for _, c := range []struct { + op, args, feature, other string + }{ + {"finetune.start", `{"model":"org/base","datasetSource":"org/data"}`, fineTune, quant}, + {"finetune.get", `{"jobId":"j1"}`, fineTune, quant}, + {"finetune.list", `{}`, fineTune, quant}, + {"finetune.checkpoints", `{"jobId":"j1"}`, fineTune, quant}, + {"quantize.start", `{"model":"org/base"}`, quant, fineTune}, + {"quantize.get", `{"jobId":"q1"}`, quant, fineTune}, + {"quantize.list", `{}`, quant, fineTune}, + } { + It(c.op+" needs its feature", func() { + none := func(string, string) bool { return false } + code, out := serveEngine(c.op, c.args, none) + Expect(code).To(Equal(http.StatusForbidden), "without any feature") + Expect(out.Error).ToNot(BeNil()) + Expect(out.Error.Code).To(Equal(contract.HostFeature), "without any feature") + otherOnly := func(_, f string) bool { return f == c.other } + code, out = serveEngine(c.op, c.args, otherOnly) + Expect(code).To(Equal(http.StatusForbidden), "with only %s", c.other) + Expect(out.Error).ToNot(BeNil()) + Expect(out.Error.Code).To(Equal(contract.HostFeature), "with only %s", c.other) + ownOnly := func(_, f string) bool { return f == c.feature } + code, out = serveEngine(c.op, c.args, ownOnly) + Expect(code).To(Equal(http.StatusOK), "with %s: %+v", c.feature, out.Error) + }) } - var got struct { - Jobs []struct { - ID string `json:"id"` - } `json:"jobs"` - } - _ = json.Unmarshal(out, &got) - var r []string - for _, j := range got.Jobs { - r = append(r, j.ID) - } - return r - } - if got := ids("u2"); len(got) != 1 || got[0] != "b" { - t.Fatalf("user 2 must see only its own job, got %v", got) - } - if got := ids("u1"); len(got) != 2 { - t.Fatalf("user 1 must see its two jobs, got %v", got) - } - // A user with no job sees none, not the others'. - if got := ids("u3"); len(got) != 0 { - t.Fatalf("user 3 must see nothing, got %v", got) - } -} + }) + + Describe("dataset source", func() { + // A dataset reference that is not an id is refused, the field is + // named, and the service is never reached. + It("must be a gallery or Hugging Face ID", func() { + bad := []string{ + "https://example.com/data.jsonl", // URL + "hf://org/data", // scheme + "file:///etc/passwd", // scheme and absolute path + "/var/data.jsonl", // absolute path + "../data", // parent segment + "org/../data", // parent segment inside an id + "org/..", // parent segment at the end + "org data", // space + "org/data/extra", // more than one separator + `C:\data`, // backslash + "", // empty + } + for _, src := range bad { + ft := &fakeFT{} + s := specs(ft) + args, _ := json.Marshal(map[string]string{"model": "org/base", "datasetSource": src}) + _, err := call(s["finetune.start"], "named-user", "u1", string(args)) + Expect(err).To(MatchError(ErrBadRequest), "%q", src) + Expect(err.Error()).To(ContainSubstring("datasetSource"), "%q", src) + Expect(ft.started).To(BeEmpty(), "%q reached the service", src) + } + }) + + // The answer the add-on process receives through the channel. + It("refusal is a 400 naming the field", func() { + code, out := serveEngine("finetune.start", `{"model":"org/base","datasetSource":"https://example.com/d.jsonl"}`, func(string, string) bool { return true }) + Expect(code).To(Equal(http.StatusBadRequest)) + Expect(out.Error).ToNot(BeNil()) + Expect(out.Error.Code).To(Equal(contract.HostBadRequest)) + Expect(out.Error.Message).To(ContainSubstring("datasetSource")) + }) + + // Keeps the ids the UI uses. + It("accepts Hugging Face and gallery IDs", func() { + for _, src := range []string{"org/data", "HuggingFaceH4/ultrachat_200k", "my-dataset", "org/data-v1.0"} { + ft := &fakeFT{} + s := specs(ft) + args, _ := json.Marshal(map[string]string{"model": "org/base", "datasetSource": src}) + _, err := call(s["finetune.start"], "named-user", "u1", string(args)) + Expect(err).ToNot(HaveOccurred(), "%q", src) + } + }) + + It("accepts a dataset ref as dataset source", func() { + ft := &fakeFT{} + ds := &fakeDatasets{} + b := New(NewInFlight(), func(string, string) bool { return true }, contract.DefaultLimits()) + Expect(RegisterEngineJobs(b, ft, nil, backends{})).To(Succeed()) + Expect(RegisterEngineDatasets(b, ds)).To(Succeed()) + raw, err := call(b.specs["finetune.dataset"], "named-user", "u1", datasetArgs("x.jsonl", "[]")) + Expect(err).ToNot(HaveOccurred()) + var out datasetResult + Expect(json.Unmarshal(raw, &out)).To(Succeed()) + start := b.specs["finetune.start"] + _, err = call(start, "named-user", "u1", `{"model":"org/base","datasetSource":"`+out.DatasetRef+`"}`) + Expect(err).ToNot(HaveOccurred()) + Expect(ft.started[len(ft.started)-1].DatasetSource).To(Equal(ds.paths[out.DatasetRef]), "the service must get the stored path") + _, err = call(start, "named-user", "u2", `{"model":"org/base","datasetSource":"`+out.DatasetRef+`"}`) + Expect(err).To(MatchError(ErrBadRequest), "another user's reference") + _, err = call(start, "named-user", "u1", `{"model":"org/base","datasetSource":"ffffffff-nope.jsonl"}`) + Expect(err).To(MatchError(ErrBadRequest), "an unknown reference") + }) + }) +}) diff --git a/core/services/addons/hostops/engine_models_test.go b/core/services/addons/hostops/engine_models_test.go index da24dc382..60cde1093 100644 --- a/core/services/addons/hostops/engine_models_test.go +++ b/core/services/addons/hostops/engine_models_test.go @@ -5,16 +5,16 @@ import ( "bytes" "context" "encoding/json" - "errors" "fmt" "net/http" "net/http/httptest" "strings" - "testing" "time" "github.com/mudler/LocalAI/core/schema" "github.com/mudler/LocalAI/pkg/addons/contract" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" ) type fakeExporter struct { @@ -68,19 +68,15 @@ func newExporter(path string) *fakeExporter { return &fakeExporter{owner: map[string]string{"j1": "u1"}, cps: oneCheckpoint, path: path} } -func modelSpecs(t *testing.T, ex Exporter, im Importer) map[string]Spec { - t.Helper() +func modelSpecs(ex Exporter, im Importer) map[string]Spec { b := New(NewInFlight(), func(string, string) bool { return true }, contract.DefaultLimits()) - if err := RegisterEngineModels(b, ex, im); err != nil { - t.Fatal(err) - } + Expect(RegisterEngineModels(b, ex, im)).To(Succeed()) return b.specs } // serveAs runs one operation through the channel handler as the given actor, // with a live envelope for it, and returns the status and the answer. -func serveAs(t *testing.T, b *Broker, op, mode, user, args string) (int, contract.HostResponse) { - t.Helper() +func serveAs(b *Broker, op, mode, user, args string) (int, contract.HostResponse) { end := b.inflight.Begin(contract.Envelope{ Request: contract.Request{AddonID: "fine-tuning", OperationID: "op1", Generation: 1, Action: "start", Arguments: json.RawMessage(`{}`)}, Authority: contract.Authority{Actor: contract.Actor{UserID: user, CredentialID: "c1", Mode: mode}, Scope: "apply", ExpiresAt: time.Now().Add(time.Minute)}, @@ -99,214 +95,175 @@ func serveAs(t *testing.T, b *Broker, op, mode, user, args string) (int, contrac // adminBroker builds a broker where every user holds the feature, and only the // user "admin" is an administrator. -func adminBroker(t *testing.T, ex Exporter, im Importer) *Broker { - t.Helper() +func adminBroker(ex Exporter, im Importer) *Broker { b := New(NewInFlight(), func(string, string) bool { return true }, contract.DefaultLimits()) b.SetAdminCheck(func(actor contract.Actor) bool { return actor.UserID == "admin" }) - if err := RegisterEngineModels(b, ex, im); err != nil { - t.Fatal(err) - } + Expect(RegisterEngineModels(b, ex, im)).To(Succeed()) return b } -func TestExportRegistersAModelAndReturnsOnlyItsName(t *testing.T) { - ex := newExporter("/models/my-model") - s := modelSpecs(t, ex, nil) - out, err := call(t, s["finetune.export"], "named-user", "u1", `{"jobId":"j1","name":"my-model","exportFormat":"lora"}`) - if err != nil { - t.Fatal(err) - } - if string(out) != `{"model":"my-model"}` { - t.Fatalf("got %s", out) - } - if strings.Contains(string(out), "/") { - t.Fatalf("result carries a path: %s", out) - } - if ex.user != "u1" || ex.req.Name != "my-model" || ex.req.ExportFormat != "lora" { - t.Fatalf("service got user=%q req=%+v", ex.user, ex.req) - } -} +var _ = Describe("Model export and import operations", func() { + It("export registers a model and returns only its name", func() { + ex := newExporter("/models/my-model") + s := modelSpecs(ex, nil) + out, err := call(s["finetune.export"], "named-user", "u1", `{"jobId":"j1","name":"my-model","exportFormat":"lora"}`) + Expect(err).ToNot(HaveOccurred()) + Expect(string(out)).To(Equal(`{"model":"my-model"}`)) + Expect(string(out)).ToNot(ContainSubstring("/"), "result carries a path") + Expect(ex.user).To(Equal("u1")) + Expect(ex.req.Name).To(Equal("my-model")) + Expect(ex.req.ExportFormat).To(Equal("lora")) + }) -func TestExportUsesTheHighestStepCheckpoint(t *testing.T) { - ex := newExporter("/models/my-model") - ex.cps = []Checkpoint{{Step: 100, Path: "/out/ckpt-100"}, {Step: 300, Path: "/out/ckpt-300"}, {Step: 200, Path: "/out/ckpt-200"}} - s := modelSpecs(t, ex, nil) - if _, err := call(t, s["finetune.export"], "named-user", "u1", `{"jobId":"j1"}`); err != nil { - t.Fatal(err) - } - if ex.req.CheckpointPath != "/out/ckpt-300" { - t.Fatalf("export got checkpoint %q, want the highest step", ex.req.CheckpointPath) - } -} + It("export uses the highest step checkpoint", func() { + ex := newExporter("/models/my-model") + ex.cps = []Checkpoint{{Step: 100, Path: "/out/ckpt-100"}, {Step: 300, Path: "/out/ckpt-300"}, {Step: 200, Path: "/out/ckpt-200"}} + s := modelSpecs(ex, nil) + _, err := call(s["finetune.export"], "named-user", "u1", `{"jobId":"j1"}`) + Expect(err).ToNot(HaveOccurred()) + Expect(ex.req.CheckpointPath).To(Equal("/out/ckpt-300"), "want the highest step") + }) -func TestExportOfAJobWithNoCheckpointIsAConflictNamingIt(t *testing.T) { - ex := newExporter("/models/my-model") - ex.cps = nil - s := modelSpecs(t, ex, nil) - _, err := call(t, s["finetune.export"], "named-user", "u1", `{"jobId":"j1"}`) - if !errors.Is(err, ErrConflict) || !strings.Contains(err.Error(), "checkpoint") { - t.Fatalf("got %v, want a conflict naming the checkpoint", err) - } - if ex.calls != 0 { - t.Fatalf("service called %d times for a job with no checkpoint", ex.calls) - } -} + It("export of a job with no checkpoint is a conflict naming it", func() { + ex := newExporter("/models/my-model") + ex.cps = nil + s := modelSpecs(ex, nil) + _, err := call(s["finetune.export"], "named-user", "u1", `{"jobId":"j1"}`) + Expect(err).To(MatchError(ErrConflict)) + Expect(err.Error()).To(ContainSubstring("checkpoint")) + Expect(ex.calls).To(Equal(0), "service called for a job with no checkpoint") + }) -func TestExportRefusesABadName(t *testing.T) { - ex := newExporter("/models/x") - s := modelSpecs(t, ex, nil) - long := strings.Repeat("a", 65) - for _, name := range []string{"../x", "a/b", " ", long} { - args, _ := json.Marshal(map[string]string{"jobId": "j1", "name": name}) - if _, err := call(t, s["finetune.export"], "named-user", "u1", string(args)); !errors.Is(err, ErrBadRequest) { - t.Fatalf("%q: want ErrBadRequest, got %v", name, err) + It("export refuses a bad name", func() { + ex := newExporter("/models/x") + s := modelSpecs(ex, nil) + long := strings.Repeat("a", 65) + for _, name := range []string{"../x", "a/b", " ", long} { + args, _ := json.Marshal(map[string]string{"jobId": "j1", "name": name}) + _, err := call(s["finetune.export"], "named-user", "u1", string(args)) + Expect(err).To(MatchError(ErrBadRequest), "%q", name) } - } - if ex.calls != 0 { - t.Fatalf("service called %d times for a refused name", ex.calls) - } -} + Expect(ex.calls).To(Equal(0), "service called for a refused name") + }) -func TestExportRefusesAnExportFormatThatIsNotListed(t *testing.T) { - ex := newExporter("/models/x") - s := modelSpecs(t, ex, nil) - for _, f := range []string{"adapter", "merged_8bit", "GGUF", "gguf "} { - args, _ := json.Marshal(map[string]string{"jobId": "j1", "exportFormat": f}) - if _, err := call(t, s["finetune.export"], "named-user", "u1", string(args)); !errors.Is(err, ErrBadRequest) { - t.Fatalf("%q: want ErrBadRequest, got %v", f, err) + It("export refuses an export format that is not listed", func() { + ex := newExporter("/models/x") + s := modelSpecs(ex, nil) + for _, f := range []string{"adapter", "merged_8bit", "GGUF", "gguf "} { + args, _ := json.Marshal(map[string]string{"jobId": "j1", "exportFormat": f}) + _, err := call(s["finetune.export"], "named-user", "u1", string(args)) + Expect(err).To(MatchError(ErrBadRequest), "%q", f) } - } - for _, f := range []string{"lora", "merged_16bit", "merged_4bit", "gguf"} { - args, _ := json.Marshal(map[string]string{"jobId": "j1", "exportFormat": f}) - if _, err := call(t, s["finetune.export"], "named-user", "u1", string(args)); err != nil { - t.Fatalf("%q: %v", f, err) + for _, f := range []string{"lora", "merged_16bit", "merged_4bit", "gguf"} { + args, _ := json.Marshal(map[string]string{"jobId": "j1", "exportFormat": f}) + _, err := call(s["finetune.export"], "named-user", "u1", string(args)) + Expect(err).ToNot(HaveOccurred(), "%q", f) } - } -} + }) -func TestExportOfAnotherUsersJobIsNotFound(t *testing.T) { - ex := newExporter("/models/x") - ex.owner = map[string]string{"j1": "u2"} - s := modelSpecs(t, ex, nil) - if _, err := call(t, s["finetune.export"], "named-user", "u1", `{"jobId":"j1"}`); !errors.Is(err, ErrNotFound) { - t.Fatalf("got %v", err) - } -} + It("export of another user's job is not found", func() { + ex := newExporter("/models/x") + ex.owner = map[string]string{"j1": "u2"} + s := modelSpecs(ex, nil) + _, err := call(s["finetune.export"], "named-user", "u1", `{"jobId":"j1"}`) + Expect(err).To(MatchError(ErrNotFound)) + }) -func TestExportRunsOncePerOperationID(t *testing.T) { - ex := newExporter("/models/my-model") - b := New(NewInFlight(), func(string, string) bool { return true }, contract.DefaultLimits()) - if err := RegisterEngineModels(b, ex, nil); err != nil { - t.Fatal(err) - } - spec := b.specs["finetune.export"] - c := Call{AddonID: "fine-tuning", OperationID: "op1", Actor: contract.Actor{UserID: "u1", Mode: "named-user"}, Arguments: json.RawMessage(`{"jobId":"j1"}`)} - first, err := b.run(context.Background(), spec, c) - if err != nil { - t.Fatal(err) - } - second, err := b.run(context.Background(), spec, c) - if err != nil { - t.Fatal(err) - } - if ex.calls != 1 || string(first) != string(second) { - t.Fatalf("calls=%d first=%s second=%s", ex.calls, first, second) - } -} + It("export runs once per operation ID", func() { + ex := newExporter("/models/my-model") + b := New(NewInFlight(), func(string, string) bool { return true }, contract.DefaultLimits()) + Expect(RegisterEngineModels(b, ex, nil)).To(Succeed()) + spec := b.specs["finetune.export"] + c := Call{AddonID: "fine-tuning", OperationID: "op1", Actor: contract.Actor{UserID: "u1", Mode: "named-user"}, Arguments: json.RawMessage(`{"jobId":"j1"}`)} + first, err := b.run(context.Background(), spec, c) + Expect(err).ToNot(HaveOccurred()) + second, err := b.run(context.Background(), spec, c) + Expect(err).ToNot(HaveOccurred()) + Expect(ex.calls).To(Equal(1)) + Expect(string(first)).To(Equal(string(second))) + }) -func TestExportRefusesAReturnedPathThatIsNotAModelName(t *testing.T) { - ex := newExporter("/models/bad name") - s := modelSpecs(t, ex, nil) - if _, err := call(t, s["finetune.export"], "named-user", "u1", `{"jobId":"j1"}`); err == nil || errors.Is(err, ErrBadRequest) { - t.Fatalf("want an outage error, got %v", err) - } -} + It("export refuses a returned path that is not a model name", func() { + ex := newExporter("/models/bad name") + s := modelSpecs(ex, nil) + _, err := call(s["finetune.export"], "named-user", "u1", `{"jobId":"j1"}`) + Expect(err).To(HaveOccurred(), "want an outage error") + Expect(err).ToNot(MatchError(ErrBadRequest), "want an outage error") + }) -func TestImportRegistersTheQuantizedModelAndReturnsOnlyItsName(t *testing.T) { - im := &fakeImporter{path: "/models/q-model"} - s := modelSpecs(t, nil, im) - out, err := call(t, s["quantize.import"], "named-user", "u1", `{"jobId":"q1","name":"q-model"}`) - if err != nil { - t.Fatal(err) - } - if string(out) != `{"model":"q-model"}` { - t.Fatalf("got %s", out) - } -} + It("import registers the quantized model and returns only its name", func() { + im := &fakeImporter{path: "/models/q-model"} + s := modelSpecs(nil, im) + out, err := call(s["quantize.import"], "named-user", "u1", `{"jobId":"q1","name":"q-model"}`) + Expect(err).ToNot(HaveOccurred()) + Expect(string(out)).To(Equal(`{"model":"q-model"}`)) + }) -func TestImportOfAJobThatIsNotCompletedIsAConflict(t *testing.T) { - im := &fakeImporter{err: fmt.Errorf("%w: job q1 is not completed", ErrConflict)} - s := modelSpecs(t, nil, im) - if _, err := call(t, s["quantize.import"], "named-user", "u1", `{"jobId":"q1"}`); !errors.Is(err, ErrConflict) { - t.Fatalf("got %v", err) - } -} + It("import of a job that is not completed is a conflict", func() { + im := &fakeImporter{err: fmt.Errorf("%w: job q1 is not completed", ErrConflict)} + s := modelSpecs(nil, im) + _, err := call(s["quantize.import"], "named-user", "u1", `{"jobId":"q1"}`) + Expect(err).To(MatchError(ErrConflict)) + }) -func TestModelOperationsAreApplyAndIdempotentAndAdminOnly(t *testing.T) { - s := modelSpecs(t, newExporter("/models/x"), &fakeImporter{}) - for _, op := range []string{"finetune.export", "quantize.import"} { - if s[op].Scope != "apply" || !s[op].Idempotent || !s[op].Admin { - t.Fatalf("%s: scope=%q idempotent=%v admin=%v", op, s[op].Scope, s[op].Idempotent, s[op].Admin) + It("model operations are apply, idempotent and admin only", func() { + s := modelSpecs(newExporter("/models/x"), &fakeImporter{}) + for _, op := range []string{"finetune.export", "quantize.import"} { + Expect(s[op].Scope).To(Equal("apply"), op) + Expect(s[op].Idempotent).To(BeTrue(), op) + Expect(s[op].Admin).To(BeTrue(), op) } - } - if s["finetune.export"].Feature != featureFineTuning || s["quantize.import"].Feature != featureQuantization { - t.Fatalf("features: %q %q", s["finetune.export"].Feature, s["quantize.import"].Feature) - } -} + Expect(s["finetune.export"].Feature).To(Equal(featureFineTuning)) + Expect(s["quantize.import"].Feature).To(Equal(featureQuantization)) + }) -func TestExportIsRefusedForAFeatureHoldingNonAdmin(t *testing.T) { - ex := newExporter("/models/my-model") - b := adminBroker(t, ex, nil) - code, out := serveAs(t, b, "finetune.export", "named-user", "u1", `{"jobId":"j1"}`) - if code != http.StatusForbidden || out.Error == nil || out.Error.Code != contract.HostFeature { - t.Fatalf("%d %+v", code, out.Error) - } - if ex.calls != 0 { - t.Fatalf("service called %d times for a non-admin", ex.calls) - } -} + It("export is refused for a feature-holding non-admin", func() { + ex := newExporter("/models/my-model") + b := adminBroker(ex, nil) + code, out := serveAs(b, "finetune.export", "named-user", "u1", `{"jobId":"j1"}`) + Expect(code).To(Equal(http.StatusForbidden)) + Expect(out.Error).ToNot(BeNil()) + Expect(out.Error.Code).To(Equal(contract.HostFeature)) + Expect(ex.calls).To(Equal(0), "service called for a non-admin") + }) -func TestExportIsAllowedForAnAdmin(t *testing.T) { - ex := newExporter("/models/my-model") - ex.owner = map[string]string{"j1": "admin"} - b := adminBroker(t, ex, nil) - code, out := serveAs(t, b, "finetune.export", "named-user", "admin", `{"jobId":"j1"}`) - if code != http.StatusOK || out.Error != nil { - t.Fatalf("%d %+v", code, out.Error) - } -} + It("export is allowed for an admin", func() { + ex := newExporter("/models/my-model") + ex.owner = map[string]string{"j1": "admin"} + b := adminBroker(ex, nil) + code, out := serveAs(b, "finetune.export", "named-user", "admin", `{"jobId":"j1"}`) + Expect(code).To(Equal(http.StatusOK)) + Expect(out.Error).To(BeNil()) + }) -func TestImportIsRefusedForAFeatureHoldingNonAdmin(t *testing.T) { - im := &fakeImporter{path: "/models/q-model"} - b := adminBroker(t, nil, im) - code, out := serveAs(t, b, "quantize.import", "named-user", "u1", `{"jobId":"q1"}`) - if code != http.StatusForbidden || out.Error == nil || out.Error.Code != contract.HostFeature { - t.Fatalf("%d %+v", code, out.Error) - } - if im.calls != 0 { - t.Fatalf("service called %d times for a non-admin", im.calls) - } -} + It("import is refused for a feature-holding non-admin", func() { + im := &fakeImporter{path: "/models/q-model"} + b := adminBroker(nil, im) + code, out := serveAs(b, "quantize.import", "named-user", "u1", `{"jobId":"q1"}`) + Expect(code).To(Equal(http.StatusForbidden)) + Expect(out.Error).ToNot(BeNil()) + Expect(out.Error.Code).To(Equal(contract.HostFeature)) + Expect(im.calls).To(Equal(0), "service called for a non-admin") + }) -func TestImportIsAllowedForAnAdmin(t *testing.T) { - im := &fakeImporter{path: "/models/q-model"} - b := adminBroker(t, nil, im) - code, out := serveAs(t, b, "quantize.import", "named-user", "admin", `{"jobId":"q1"}`) - if code != http.StatusOK || out.Error != nil { - t.Fatalf("%d %+v", code, out.Error) - } -} + It("import is allowed for an admin", func() { + im := &fakeImporter{path: "/models/q-model"} + b := adminBroker(nil, im) + code, out := serveAs(b, "quantize.import", "named-user", "admin", `{"jobId":"q1"}`) + Expect(code).To(Equal(http.StatusOK)) + Expect(out.Error).To(BeNil()) + }) -// The broker does not shortcut the local owner: the admin check decides in -// every mode, so a local-owner actor the check refuses is refused. -func TestTheAdminCheckDecidesInEveryMode(t *testing.T) { - ex := newExporter("/models/my-model") - ex.owner = map[string]string{"j1": "u1"} - b := adminBroker(t, ex, nil) - if code, out := serveAs(t, b, "finetune.export", "local-owner", "u1", `{"jobId":"j1"}`); code != http.StatusForbidden || out.Error.Code != contract.HostFeature { - t.Fatalf("local-owner actor the check refuses: %d %+v", code, out.Error) - } - if ex.calls != 0 { - t.Fatalf("service called %d times for a refused actor", ex.calls) - } -} + // The broker does not shortcut the local owner: the admin check decides in + // every mode, so a local-owner actor the check refuses is refused. + It("lets the admin check decide in every mode", func() { + ex := newExporter("/models/my-model") + ex.owner = map[string]string{"j1": "u1"} + b := adminBroker(ex, nil) + code, out := serveAs(b, "finetune.export", "local-owner", "u1", `{"jobId":"j1"}`) + Expect(code).To(Equal(http.StatusForbidden), "local-owner actor the check refuses") + Expect(out.Error).ToNot(BeNil()) + Expect(out.Error.Code).To(Equal(contract.HostFeature)) + Expect(ex.calls).To(Equal(0), "service called for a refused actor") + }) +}) diff --git a/core/services/addons/hostops/engine_progress_test.go b/core/services/addons/hostops/engine_progress_test.go index 542174cdd..914d984fb 100644 --- a/core/services/addons/hostops/engine_progress_test.go +++ b/core/services/addons/hostops/engine_progress_test.go @@ -5,14 +5,14 @@ import ( "bytes" "context" "encoding/json" - "errors" "net/http" "net/http/httptest" "sync" - "testing" "time" "github.com/mudler/LocalAI/pkg/addons/contract" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" ) // fakeSource stands in for a job service. Events are fed by the test through @@ -76,160 +76,125 @@ func (f *fakeSource) stopCount() int { return len(f.stops) } -func registerProgress(t *testing.T, src ProgressSource) (*Broker, map[string]Spec) { - t.Helper() +func registerProgress(src ProgressSource) (*Broker, map[string]Spec) { b := New(NewInFlight(), func(string, string) bool { return true }, contract.DefaultLimits()) p := NewProgresser(src) - if err := RegisterEngineProgress(b, p, p); err != nil { - t.Fatal(err) - } + Expect(RegisterEngineProgress(b, p, p)).To(Succeed()) return b, b.specs } -func TestProgressReturnsTheNextEventWithinTheWait(t *testing.T) { - src := newFakeSource() - _, s := registerProgress(t, src) - src.events <- Progress{JobID: "j1", Status: "training", Percent: 40} - out, err := call(t, s["finetune.progress"], "named-user", "u1", `{"jobId":"j1","waitSeconds":2}`) - if err != nil { - t.Fatal(err) - } - var got struct { - Percent float32 `json:"percent"` - Done bool `json:"done"` - } - _ = json.Unmarshal(out, &got) - if got.Percent != 40 || got.Done { - t.Fatalf("want percent 40 and not done: %s", out) - } -} +var _ = Describe("Progress and stop operations", func() { + var ( + src *fakeSource + b *Broker + s map[string]Spec + ) -func TestProgressWithNoEventReturnsTheSnapshotWhenTheWaitEnds(t *testing.T) { - src := newFakeSource() - _, s := registerProgress(t, src) - start := time.Now() - out, err := call(t, s["finetune.progress"], "named-user", "u1", `{"jobId":"j1","waitSeconds":1}`) - elapsed := time.Since(start) - if err != nil { - t.Fatal(err) - } - if elapsed < 900*time.Millisecond || elapsed >= 3*time.Second { - t.Fatalf("the call should wait about one second, took %s", elapsed) - } - var got struct { - Status string `json:"status"` - Done bool `json:"done"` - } - _ = json.Unmarshal(out, &got) - if got.Status != "training" || got.Done { - t.Fatalf("want the current status of the job: %s", out) - } -} + BeforeEach(func() { + src = newFakeSource() + b, s = registerProgress(src) + }) -func TestProgressWaitIsCapped(t *testing.T) { - src := newFakeSource() - _, s := registerProgress(t, src) - src.events <- Progress{JobID: "j1", Status: "training"} - if _, err := call(t, s["finetune.progress"], "named-user", "u1", `{"jobId":"j1","waitSeconds":600}`); err != nil { - t.Fatal(err) - } - src.mu.Lock() - left := time.Until(src.deadline) - src.mu.Unlock() - if left <= 24*time.Second || left > 25*time.Second+500*time.Millisecond { - t.Fatalf("the stream should run under a wait of 25 s, got about %s", left) - } -} - -func TestProgressReportsDoneForAFinalStatus(t *testing.T) { - src := newFakeSource() - _, s := registerProgress(t, src) - src.events <- Progress{JobID: "j1", Status: "completed"} - out, err := call(t, s["finetune.progress"], "named-user", "u1", `{"jobId":"j1","waitSeconds":2}`) - if err != nil { - t.Fatal(err) - } - var got struct { - Done bool `json:"done"` - } - _ = json.Unmarshal(out, &got) - if !got.Done { - t.Fatalf("a completed job is done: %s", out) - } -} - -func TestProgressOfAnotherUsersJobIsNotFound(t *testing.T) { - src := newFakeSource() - _, s := registerProgress(t, src) - if _, err := call(t, s["finetune.progress"], "named-user", "u2", `{"jobId":"j1","waitSeconds":1}`); !errors.Is(err, ErrNotFound) { - t.Fatalf("got %v", err) - } - if _, err := call(t, s["quantize.progress"], "named-user", "u2", `{"jobId":"q1","waitSeconds":1}`); !errors.Is(err, ErrNotFound) { - t.Fatalf("quantize: got %v", err) - } -} - -func TestStopIsIdempotentPerOperation(t *testing.T) { - src := newFakeSource() - b, s := registerProgress(t, src) - c := Call{AddonID: "fine-tuning", OperationID: "op1", Actor: contract.Actor{UserID: "u1", Mode: "named-user"}, Arguments: json.RawMessage(`{"jobId":"j1"}`)} - for range 2 { - if _, err := b.run(context.Background(), s["finetune.stop"], c); err != nil { - t.Fatal(err) + It("returns the next event within the wait", func() { + src.events <- Progress{JobID: "j1", Status: "training", Percent: 40} + out, err := call(s["finetune.progress"], "named-user", "u1", `{"jobId":"j1","waitSeconds":2}`) + Expect(err).ToNot(HaveOccurred()) + var got struct { + Percent float32 `json:"percent"` + Done bool `json:"done"` } - } - if n := src.stopCount(); n != 1 { - t.Fatalf("two calls with one operation ID must stop once, stopped %d times", n) - } -} + _ = json.Unmarshal(out, &got) + Expect(got.Percent).To(BeEquivalentTo(40), "want percent 40: %s", out) + Expect(got.Done).To(BeFalse(), "want not done: %s", out) + }) -func TestStopKeepsTheCheckpointOnlyWhenAsked(t *testing.T) { - src := newFakeSource() - _, s := registerProgress(t, src) - if _, err := call(t, s["finetune.stop"], "named-user", "u1", `{"jobId":"j1","saveCheckpoint":true}`); err != nil { - t.Fatal(err) - } - if _, err := call(t, s["finetune.stop"], "named-user", "u1", `{"jobId":"j1"}`); err != nil { - t.Fatal(err) - } - if len(src.stops) != 2 || !src.stops[0] || src.stops[1] { - t.Fatalf("saveCheckpoint must reach StopJob only when true: %v", src.stops) - } -} + It("returns the snapshot when the wait ends with no event", func() { + start := time.Now() + out, err := call(s["finetune.progress"], "named-user", "u1", `{"jobId":"j1","waitSeconds":1}`) + elapsed := time.Since(start) + Expect(err).ToNot(HaveOccurred()) + Expect(elapsed).To(BeNumerically(">=", 900*time.Millisecond), "the call should wait about one second") + Expect(elapsed).To(BeNumerically("<", 3*time.Second), "the call should wait about one second") + var got struct { + Status string `json:"status"` + Done bool `json:"done"` + } + _ = json.Unmarshal(out, &got) + Expect(got.Status).To(Equal("training"), "want the current status of the job: %s", out) + Expect(got.Done).To(BeFalse(), "want the current status of the job: %s", out) + }) -func TestStopOfAnotherUsersJobIsNotFound(t *testing.T) { - src := newFakeSource() - _, s := registerProgress(t, src) - if _, err := call(t, s["quantize.stop"], "named-user", "u2", `{"jobId":"q1"}`); !errors.Is(err, ErrNotFound) { - t.Fatalf("got %v", err) - } - if src.stopCount() != 0 { - t.Fatal("another user's job must not be stopped") - } -} + It("caps the wait", func() { + src.events <- Progress{JobID: "j1", Status: "training"} + _, err := call(s["finetune.progress"], "named-user", "u1", `{"jobId":"j1","waitSeconds":600}`) + Expect(err).ToNot(HaveOccurred()) + src.mu.Lock() + left := time.Until(src.deadline) + src.mu.Unlock() + Expect(left).To(BeNumerically(">", 24*time.Second), "the stream should run under a wait of 25 s") + Expect(left).To(BeNumerically("<=", 25*time.Second+500*time.Millisecond), "the stream should run under a wait of 25 s") + }) -func TestACancelledContextEndsTheWait(t *testing.T) { - src := newFakeSource() - b, _ := registerProgress(t, src) - in := b.inflight - end := in.Begin(env("fine-tuning", "op1", "read", time.Now().Add(time.Minute))) - defer end() - h := b.Handler("fine-tuning", "secret", []string{"finetune."}) + It("reports done for a final status", func() { + src.events <- Progress{JobID: "j1", Status: "completed"} + out, err := call(s["finetune.progress"], "named-user", "u1", `{"jobId":"j1","waitSeconds":2}`) + Expect(err).ToNot(HaveOccurred()) + var got struct { + Done bool `json:"done"` + } + _ = json.Unmarshal(out, &got) + Expect(got.Done).To(BeTrue(), "a completed job is done: %s", out) + }) - body, _ := json.Marshal(contract.HostRequest{OperationID: "op1", Op: "finetune.progress", Arguments: json.RawMessage(`{"jobId":"j1","waitSeconds":25}`)}) - ctx, cancel := context.WithCancel(context.Background()) - cancel() - req := httptest.NewRequest(http.MethodPost, "/host", bytes.NewReader(body)).WithContext(ctx) - req.Header.Set("X-Addon-Host-Credential", "secret") - rec := httptest.NewRecorder() - start := time.Now() - h.ServeHTTP(rec, req) - if d := time.Since(start); d > time.Second { - t.Fatalf("a cancelled call must end at once, took %s", d) - } - var out contract.HostResponse - _ = json.Unmarshal(rec.Body.Bytes(), &out) - if rec.Code != http.StatusServiceUnavailable || out.Error == nil || out.Error.Code != contract.HostUnavailable { - t.Fatalf("want 503 unavailable, got %d %+v", rec.Code, out.Error) - } -} + It("does not find the progress of another user's job", func() { + _, err := call(s["finetune.progress"], "named-user", "u2", `{"jobId":"j1","waitSeconds":1}`) + Expect(err).To(MatchError(ErrNotFound)) + _, err = call(s["quantize.progress"], "named-user", "u2", `{"jobId":"q1","waitSeconds":1}`) + Expect(err).To(MatchError(ErrNotFound), "quantize") + }) + + It("stops once per operation", func() { + c := Call{AddonID: "fine-tuning", OperationID: "op1", Actor: contract.Actor{UserID: "u1", Mode: "named-user"}, Arguments: json.RawMessage(`{"jobId":"j1"}`)} + for range 2 { + _, err := b.run(context.Background(), s["finetune.stop"], c) + Expect(err).ToNot(HaveOccurred()) + } + Expect(src.stopCount()).To(Equal(1), "two calls with one operation ID must stop once") + }) + + It("keeps the checkpoint only when asked", func() { + _, err := call(s["finetune.stop"], "named-user", "u1", `{"jobId":"j1","saveCheckpoint":true}`) + Expect(err).ToNot(HaveOccurred()) + _, err = call(s["finetune.stop"], "named-user", "u1", `{"jobId":"j1"}`) + Expect(err).ToNot(HaveOccurred()) + Expect(src.stops).To(Equal([]bool{true, false}), "saveCheckpoint must reach StopJob only when true") + }) + + It("does not find the stop of another user's job", func() { + _, err := call(s["quantize.stop"], "named-user", "u2", `{"jobId":"q1"}`) + Expect(err).To(MatchError(ErrNotFound)) + Expect(src.stopCount()).To(Equal(0), "another user's job must not be stopped") + }) + + It("ends the wait when the context is cancelled", func() { + in := b.inflight + end := in.Begin(env("fine-tuning", "op1", "read", time.Now().Add(time.Minute))) + DeferCleanup(end) + h := b.Handler("fine-tuning", "secret", []string{"finetune."}) + + body, _ := json.Marshal(contract.HostRequest{OperationID: "op1", Op: "finetune.progress", Arguments: json.RawMessage(`{"jobId":"j1","waitSeconds":25}`)}) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + req := httptest.NewRequest(http.MethodPost, "/host", bytes.NewReader(body)).WithContext(ctx) + req.Header.Set("X-Addon-Host-Credential", "secret") + rec := httptest.NewRecorder() + start := time.Now() + h.ServeHTTP(rec, req) + Expect(time.Since(start)).To(BeNumerically("<=", time.Second), "a cancelled call must end at once") + var out contract.HostResponse + _ = json.Unmarshal(rec.Body.Bytes(), &out) + Expect(rec.Code).To(Equal(http.StatusServiceUnavailable)) + Expect(out.Error).ToNot(BeNil()) + Expect(out.Error.Code).To(Equal(contract.HostUnavailable)) + }) +}) diff --git a/core/services/addons/hostops/registry_test.go b/core/services/addons/hostops/registry_test.go index eea1612a0..609e2a34a 100644 --- a/core/services/addons/hostops/registry_test.go +++ b/core/services/addons/hostops/registry_test.go @@ -5,81 +5,80 @@ import ( "context" "encoding/json" "net/http" - "testing" "time" "github.com/mudler/LocalAI/pkg/addons/contract" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" ) func noop(context.Context, Call) (json.RawMessage, error) { return nil, nil } -func TestRegisterRefusesWhatIsMalformed(t *testing.T) { - b := New(NewInFlight(), nil, contract.DefaultLimits()) - if err := b.Register(Spec{Op: "finetune.get", Scope: "read", Feature: "fine_tuning", Handler: noop}); err != nil { - t.Fatal(err) - } - for name, s := range map[string]Spec{ - "duplicate": {Op: "finetune.get", Scope: "read", Feature: "fine_tuning", Handler: noop}, - "bad name": {Op: "Finetune", Scope: "read", Feature: "fine_tuning", Handler: noop}, - "bad scope": {Op: "finetune.x", Scope: "preview", Feature: "fine_tuning", Handler: noop}, - "no feature": {Op: "finetune.y", Scope: "read", Handler: noop}, - "no handler": {Op: "finetune.z", Scope: "read", Feature: "fine_tuning"}, - } { - if b.Register(s) == nil { - t.Errorf("%s: accepted", name) +var _ = Describe("Operation registry", func() { + It("refuses what is malformed", func() { + b := New(NewInFlight(), nil, contract.DefaultLimits()) + Expect(b.Register(Spec{Op: "finetune.get", Scope: "read", Feature: "fine_tuning", Handler: noop})).To(Succeed()) + for name, s := range map[string]Spec{ + "duplicate": {Op: "finetune.get", Scope: "read", Feature: "fine_tuning", Handler: noop}, + "bad name": {Op: "Finetune", Scope: "read", Feature: "fine_tuning", Handler: noop}, + "bad scope": {Op: "finetune.x", Scope: "preview", Feature: "fine_tuning", Handler: noop}, + "no feature": {Op: "finetune.y", Scope: "read", Handler: noop}, + "no handler": {Op: "finetune.z", Scope: "read", Feature: "fine_tuning"}, + } { + Expect(b.Register(s)).ToNot(Succeed(), "%s: accepted", name) } - } -} + }) -func TestAnUnregisteredOperationUnderAnAllowedPrefixIsRefused(t *testing.T) { - f := newFixture(t, always) - end := f.inflight.Begin(env("fine-tuning", "op1", "apply", time.Now().Add(time.Minute))) - defer end() - if code, out := f.call("secret", "op1", "finetune.nothing", `{}`); code != http.StatusForbidden || out.Error.Code != "unknown_op" { - t.Fatalf("%d %+v", code, out.Error) - } -} + It("refuses an unregistered operation under an allowed prefix", func() { + f := newFixture(always) + end := f.inflight.Begin(env("fine-tuning", "op1", "apply", time.Now().Add(time.Minute))) + DeferCleanup(end) + code, out := f.call("secret", "op1", "finetune.nothing", `{}`) + Expect(code).To(Equal(http.StatusForbidden)) + Expect(out.Error).ToNot(BeNil()) + Expect(out.Error.Code).To(Equal("unknown_op")) + }) -func TestABrokerWithoutAFeatureCheckRefusesEverything(t *testing.T) { - f := newFixture(t, nil) - end := f.inflight.Begin(env("fine-tuning", "op1", "read", time.Now().Add(time.Minute))) - defer end() - if code, out := f.call("secret", "op1", "finetune.get", `{}`); code != http.StatusForbidden || out.Error.Code != "feature" { - t.Fatalf("%d %+v", code, out.Error) - } -} + It("refuses everything for a broker without a feature check", func() { + f := newFixture(nil) + end := f.inflight.Begin(env("fine-tuning", "op1", "read", time.Now().Add(time.Minute))) + DeferCleanup(end) + code, out := f.call("secret", "op1", "finetune.get", `{}`) + Expect(code).To(Equal(http.StatusForbidden)) + Expect(out.Error).ToNot(BeNil()) + Expect(out.Error.Code).To(Equal("feature")) + }) -func TestTheLocalOwnerNeedsNoFeatureCheck(t *testing.T) { - f := newFixture(t, func(string, string) bool { return false }) - e := env("fine-tuning", "op1", "read", time.Now().Add(time.Minute)) - e.Authority.Actor.Mode = "local-owner" - end := f.inflight.Begin(e) - defer end() - if code, _ := f.call("secret", "op1", "finetune.get", `{}`); code != http.StatusOK { - t.Fatalf("%d", code) - } -} + It("needs no feature check for the local owner", func() { + f := newFixture(func(string, string) bool { return false }) + e := env("fine-tuning", "op1", "read", time.Now().Add(time.Minute)) + e.Authority.Actor.Mode = "local-owner" + end := f.inflight.Begin(e) + DeferCleanup(end) + code, _ := f.call("secret", "op1", "finetune.get", `{}`) + Expect(code).To(Equal(http.StatusOK)) + }) -func TestAnotherAddonCannotUseTheOperationOfALiveRequest(t *testing.T) { - f := newFixture(t, always) - end := f.inflight.Begin(env("quantization", "op1", "read", time.Now().Add(time.Minute))) - defer end() - if code, out := f.call("secret", "op1", "finetune.get", `{}`); code != http.StatusForbidden || out.Error.Code != "no_live_operation" { - t.Fatalf("%d %+v", code, out.Error) - } -} + It("does not let another add-on use the operation of a live request", func() { + f := newFixture(always) + end := f.inflight.Begin(env("quantization", "op1", "read", time.Now().Add(time.Minute))) + DeferCleanup(end) + code, out := f.call("secret", "op1", "finetune.get", `{}`) + Expect(code).To(Equal(http.StatusForbidden)) + Expect(out.Error).ToNot(BeNil()) + Expect(out.Error.Code).To(Equal("no_live_operation")) + }) -func TestTheEnvelopeIsGoneWhenTheRequestEnds(t *testing.T) { - in := NewInFlight() - e := env("fine-tuning", "op1", "read", time.Now().Add(time.Minute)) - end1, end2 := in.Begin(e), in.Begin(e) - end1() - end1() - if _, ok := in.Lookup("fine-tuning", "op1"); !ok { - t.Fatal("the second request still runs") - } - end2() - if _, ok := in.Lookup("fine-tuning", "op1"); ok { - t.Fatal("the envelope stayed") - } -} + It("drops the envelope when the request ends", func() { + in := NewInFlight() + e := env("fine-tuning", "op1", "read", time.Now().Add(time.Minute)) + end1, end2 := in.Begin(e), in.Begin(e) + end1() + end1() + _, ok := in.Lookup("fine-tuning", "op1") + Expect(ok).To(BeTrue(), "the second request still runs") + end2() + _, ok = in.Lookup("fine-tuning", "op1") + Expect(ok).To(BeFalse(), "the envelope stayed") + }) +}) diff --git a/core/services/addons/hostops/suite_test.go b/core/services/addons/hostops/suite_test.go new file mode 100644 index 000000000..70cc5fcff --- /dev/null +++ b/core/services/addons/hostops/suite_test.go @@ -0,0 +1,14 @@ +// SPDX-License-Identifier: MIT +package hostops + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestHostops(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Add-on host operations suite") +}