Merge 1b054d7b2a into 160660e572

llm/llm_windows.go

@@ -10,6 +10,22 @@ const (
 	CREATE_NO_WINDOW            = 0x08000000
 )
 
+var (
+	kernel32           = syscall.NewLazyDLL("kernel32.dll")
+	setThreadExecState = kernel32.NewProc("SetThreadExecutionState")
+)
+
+// preventSleep keeps the system awake for the duration of a long operation.
+// Call the returned function (or defer it) to restore normal sleep behavior.
+func preventSleep() func() {
+	const esContinuous     = uintptr(0x80000000)
+	const esSystemRequired = uintptr(0x00000001)
+	setThreadExecState.Call(esContinuous | esSystemRequired) //nolint:errcheck
+	return func() {
+		setThreadExecState.Call(esContinuous) //nolint:errcheck
+	}
+}
+
 var LlamaServerSysProcAttr = &syscall.SysProcAttr{
 	// Wire up the default error handling logic If for some reason a DLL is
 	// missing in the path this will pop up a GUI Dialog explaining the fault so

llm/server.go

@@ -1586,6 +1586,7 @@ func (s *llmServer) Completion(ctx context.Context, req CompletionRequest, fn fu
 		return err
 	}
 	defer s.sem.Release(1)
+	defer preventSleep()()
 
 	// put an upper limit on num_predict to avoid the model running on forever
 	if req.Options.NumPredict < 0 || req.Options.NumPredict > 10*s.options.NumCtx {

llm/sleep_common.go

@@ -0,0 +1,7 @@
+//go:build !windows
+
+package llm
+
+func preventSleep() func() {
+	return func() {}
+}

server/download.go

@@ -215,6 +215,7 @@ func newBackoff(maxBackoff time.Duration) func(ctx context.Context) error {
 
 func (b *blobDownload) run(ctx context.Context, requestURL *url.URL, opts *registryOptions) error {
 	defer blobDownloadManager.Delete(b.Digest)
+	defer preventSleep()()
 	ctx, b.CancelFunc = context.WithCancel(ctx)
 
 	file, err := os.OpenFile(b.Name+"-partial", os.O_CREATE|os.O_RDWR, 0o644)

server/sleep_common.go

@@ -0,0 +1,7 @@
+//go:build !windows
+
+package server
+
+func preventSleep() func() {
+	return func() {}
+}

server/sleep_windows.go

@@ -0,0 +1,12 @@
+package server
+
+import "golang.org/x/sys/windows"
+
+// preventSleep keeps the system awake for the duration of a long operation.
+// Call the returned function (or defer it) to restore normal sleep behavior.
+func preventSleep() func() {
+	windows.SetThreadExecutionState(windows.ES_CONTINUOUS | windows.ES_SYSTEM_REQUIRED) //nolint:errcheck
+	return func() {
+		windows.SetThreadExecutionState(windows.ES_CONTINUOUS) //nolint:errcheck
+	}
+}