|
| 1 | +package application |
| 2 | + |
| 3 | +import ( |
| 4 | + "context" |
| 5 | + "os" |
| 6 | + "path/filepath" |
| 7 | + |
| 8 | + . "github.com/onsi/ginkgo/v2" |
| 9 | + . "github.com/onsi/gomega" |
| 10 | + |
| 11 | + "github.com/mudler/LocalAI/core/config" |
| 12 | + "github.com/mudler/LocalAI/pkg/modelartifacts" |
| 13 | + "github.com/mudler/LocalAI/pkg/system" |
| 14 | +) |
| 15 | + |
| 16 | +// blockingMaterializer parks inside Ensure until release is closed, so a spec |
| 17 | +// can hold the startup preload open and observe readiness while it is still |
| 18 | +// running. This mirrors how a real managed-artifact preload behaves: since |
| 19 | +// #10949 it downloads a HuggingFace snapshot, which for a large model is tens |
| 20 | +// of GB and can take a long time. |
| 21 | +type blockingMaterializer struct { |
| 22 | + seen chan struct{} |
| 23 | + release chan struct{} |
| 24 | +} |
| 25 | + |
| 26 | +func (b *blockingMaterializer) Ensure(ctx context.Context, _ string, spec modelartifacts.Spec) (modelartifacts.Result, error) { |
| 27 | + select { |
| 28 | + case b.seen <- struct{}{}: |
| 29 | + default: |
| 30 | + } |
| 31 | + select { |
| 32 | + case <-b.release: |
| 33 | + case <-ctx.Done(): |
| 34 | + return modelartifacts.Result{}, ctx.Err() |
| 35 | + } |
| 36 | + spec.Resolved = &modelartifacts.Resolved{ |
| 37 | + Endpoint: "https://huggingface.co", |
| 38 | + Revision: "0123456789abcdef0123456789abcdef01234567", |
| 39 | + CacheKey: "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", |
| 40 | + } |
| 41 | + return modelartifacts.Result{Spec: spec}, nil |
| 42 | +} |
| 43 | + |
| 44 | +var _ = Describe("application readiness", func() { |
| 45 | + var modelsPath string |
| 46 | + var state *system.SystemState |
| 47 | + |
| 48 | + BeforeEach(func() { |
| 49 | + modelsPath = GinkgoT().TempDir() |
| 50 | + var err error |
| 51 | + state, err = system.GetSystemState( |
| 52 | + system.WithModelPath(modelsPath), |
| 53 | + system.WithBackendPath(GinkgoT().TempDir()), |
| 54 | + ) |
| 55 | + Expect(err).NotTo(HaveOccurred()) |
| 56 | + }) |
| 57 | + |
| 58 | + It("is not ready before the startup sequence has run", func() { |
| 59 | + app := newApplication(&config.ApplicationConfig{ |
| 60 | + Context: context.Background(), |
| 61 | + SystemState: state, |
| 62 | + }) |
| 63 | + |
| 64 | + // A constructed-but-unstarted application must never advertise itself |
| 65 | + // as able to serve traffic: /readyz reads this, and a load balancer |
| 66 | + // that believes it will route requests to a process that cannot answer. |
| 67 | + Expect(app.Ready()).To(BeFalse()) |
| 68 | + }) |
| 69 | + |
| 70 | + It("stays not-ready while the startup model preload is still running", func() { |
| 71 | + blocker := &blockingMaterializer{ |
| 72 | + seen: make(chan struct{}, 1), |
| 73 | + release: make(chan struct{}), |
| 74 | + } |
| 75 | + app := newApplication(&config.ApplicationConfig{ |
| 76 | + Context: context.Background(), |
| 77 | + SystemState: state, |
| 78 | + ModelArtifactMaterializer: blocker, |
| 79 | + }) |
| 80 | + |
| 81 | + Expect(os.WriteFile(filepath.Join(modelsPath, "managed.yaml"), []byte(` |
| 82 | +name: managed |
| 83 | +backend: transformers |
| 84 | +artifacts: |
| 85 | + - source: {type: huggingface, repo: owner/repo} |
| 86 | +parameters: {model: owner/repo} |
| 87 | +`), 0644)).To(Succeed()) |
| 88 | + Expect(app.ModelConfigLoader().LoadModelConfigsFromPath(modelsPath)).To(Succeed()) |
| 89 | + |
| 90 | + preloadDone := make(chan error, 1) |
| 91 | + go func() { |
| 92 | + preloadDone <- app.ModelConfigLoader().PreloadWithContext(context.Background(), modelsPath) |
| 93 | + }() |
| 94 | + |
| 95 | + // Preload is now parked inside artifact materialization — exactly the |
| 96 | + // window in which a Kubernetes Service would otherwise send this |
| 97 | + // replica half the cluster's traffic. |
| 98 | + Eventually(blocker.seen).Should(Receive()) |
| 99 | + Consistently(app.Ready).Should(BeFalse()) |
| 100 | + |
| 101 | + close(blocker.release) |
| 102 | + Expect(<-preloadDone).NotTo(HaveOccurred()) |
| 103 | + |
| 104 | + // Preload finished, but the rest of the startup sequence has not, so |
| 105 | + // the application is still not ready. |
| 106 | + Expect(app.Ready()).To(BeFalse()) |
| 107 | + }) |
| 108 | + |
| 109 | + It("becomes ready once the startup sequence completes", func() { |
| 110 | + ctx, cancel := context.WithCancel(context.Background()) |
| 111 | + DeferCleanup(cancel) |
| 112 | + |
| 113 | + app, err := New( |
| 114 | + config.WithContext(ctx), |
| 115 | + config.WithSystemState(state), |
| 116 | + ) |
| 117 | + Expect(err).NotTo(HaveOccurred()) |
| 118 | + DeferCleanup(func() { _ = app.Shutdown() }) |
| 119 | + |
| 120 | + Expect(app.Ready()).To(BeTrue()) |
| 121 | + }) |
| 122 | +}) |
0 commit comments