diff --git a/Makefile b/Makefile index cf3248c3b..7db3ea11e 100644 --- a/Makefile +++ b/Makefile @@ -240,7 +240,7 @@ test-ci-scripts: ## pure stdlib on purpose so they run without any backend venv; the list is ## explicit because their siblings (model_identity_test) import grpc and the ## generated protobufs, which only exist inside a built backend. -PYTHON_HELPER_TESTS?=python_utils_test vllm_utils_test model_utils_test mlx_utils_test parent_watch_test +PYTHON_HELPER_TESTS?=python_utils_test vllm_utils_test model_utils_test mlx_utils_test parent_watch_test temp_utils_test test-python-helpers: cd backend/python/common && python3 -m unittest $(PYTHON_HELPER_TESTS) diff --git a/backend/go/crispasr/gocrispasr.go b/backend/go/crispasr/gocrispasr.go index be431165d..4f5b91d6a 100644 --- a/backend/go/crispasr/gocrispasr.go +++ b/backend/go/crispasr/gocrispasr.go @@ -615,10 +615,10 @@ func (w *CrispASR) TTSStream(req *pb.TTSRequest, results chan []byte) error { return fmt.Errorf("crispasr: tempfile: %w", err) } dst := tmp.Name() + defer func() { _ = os.Remove(dst) }() if err := tmp.Close(); err != nil { return fmt.Errorf("crispasr: close tempfile: %w", err) } - defer func() { _ = os.Remove(dst) }() if err := writeWAV(dst, pcm, w.sampleRate); err != nil { return err diff --git a/backend/go/stablediffusion-ggml/cpp/gosd.cpp b/backend/go/stablediffusion-ggml/cpp/gosd.cpp index 7722e8d06..b876df256 100644 --- a/backend/go/stablediffusion-ggml/cpp/gosd.cpp +++ b/backend/go/stablediffusion-ggml/cpp/gosd.cpp @@ -1144,17 +1144,25 @@ static uint8_t* load_and_resize_image(const char* path, int target_width, int ta // Write sd.cpp's audio buffer to a temp WAV file (IEEE float, interleaved). // sd_audio_t.data is planar (all channel 0 samples, then channel 1, etc.) — we // interleave on the fly so ffmpeg's standard wav demuxer can read it directly. -// Returns 0 on success and fills wav_path (must be at least 64 bytes). +// Returns 0 on success and fills wav_path. static int write_planar_float_wav(const sd_audio_t* a, char* wav_path, size_t wav_path_sz) { if (!a || !a->data || a->sample_count == 0 || a->channels == 0 || a->sample_rate == 0) { return -1; } - snprintf(wav_path, wav_path_sz, "/tmp/gosd-audio-XXXXXX.wav"); + const char* temp_dir = getenv("TMPDIR"); + if (!temp_dir || temp_dir[0] == '\0') { + temp_dir = "/tmp"; + } + int path_len = snprintf(wav_path, wav_path_sz, "%s/gosd-audio-XXXXXX.wav", temp_dir); + if (path_len < 0 || (size_t)path_len >= wav_path_sz) { + fprintf(stderr, "temporary directory path is too long\n"); + return -1; + } int fd = mkstemps(wav_path, 4); if (fd < 0) { perror("mkstemps wav"); return -1; } FILE* f = fdopen(fd, "wb"); - if (!f) { perror("fdopen wav"); close(fd); return -1; } + if (!f) { perror("fdopen wav"); close(fd); unlink(wav_path); return -1; } uint64_t frames = a->sample_count; uint32_t channels = a->channels; @@ -1221,7 +1229,7 @@ static int ffmpeg_mux_raw_to_mp4(sd_image_t* frames, int num_frames, int fps, snprintf(fps_str, sizeof(fps_str), "%d", fps); // Optional audio: write a temp WAV file if the model produced audio. - char wav_path[64] = {0}; + char wav_path[4096] = {0}; bool have_audio = false; if (audio && audio->data && audio->sample_count > 0 && audio->channels > 0 && audio->sample_rate > 0) { if (write_planar_float_wav(audio, wav_path, sizeof(wav_path)) == 0) { diff --git a/backend/python/chatterbox/backend.py b/backend/python/chatterbox/backend.py index 016925806..996d4e50b 100644 --- a/backend/python/chatterbox/backend.py +++ b/backend/python/chatterbox/backend.py @@ -19,6 +19,7 @@ import grpc sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'common')) sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'common')) from grpc_auth import get_auth_interceptors +from temp_utils import cleanup_paths import tempfile @@ -115,11 +116,6 @@ def merge_audio_files(audio_files, output_path, sample_rate): # Save the merged audio ta.save(output_path, merged_waveform, sample_rate) - # Clean up temporary files - for audio_file in audio_files: - if os.path.exists(audio_file): - os.remove(audio_file) - _ONE_DAY_IN_SECONDS = 60 * 60 * 24 # If MAX_WORKERS are specified in the environment use it, otherwise default to 1 @@ -226,19 +222,20 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): text_chunks = split_text_at_word_boundary(request.text, max_length=250) print(f"Splitting text into chunks of 250 characters: {len(text_chunks)}", file=sys.stderr) # Generate audio for each chunk - temp_audio_files = [] - for i, chunk in enumerate(text_chunks): - # Generate audio for this chunk - wav = self.model.generate(chunk, **kwargs) - - # Create temporary file for this chunk - temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.wav') - temp_file.close() - ta.save(temp_file.name, wav, self.model.sr) - temp_audio_files.append(temp_file.name) - - # Merge all audio files - merge_audio_files(temp_audio_files, request.dst, self.model.sr) + with cleanup_paths() as temp_audio_files: + for i, chunk in enumerate(text_chunks): + # Generate audio for this chunk + wav = self.model.generate(chunk, **kwargs) + + # Register ownership before saving so a partial write is + # removed too when generation or encoding fails. + temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.wav') + temp_file.close() + temp_audio_files.append(temp_file.name) + ta.save(temp_file.name, wav, self.model.sr) + + # Merge all audio files + merge_audio_files(temp_audio_files, request.dst, self.model.sr) else: # Generate audio using ChatterboxTTS for short text wav = self.model.generate(request.text, **kwargs) diff --git a/backend/python/common/temp_utils.py b/backend/python/common/temp_utils.py new file mode 100644 index 000000000..e67f96ff9 --- /dev/null +++ b/backend/python/common/temp_utils.py @@ -0,0 +1,36 @@ +import base64 +import contextlib +import os +import tempfile + + +@contextlib.contextmanager +def materialize_base64(data, suffix=""): + """Materialize base64 data for a path-only library and always remove it.""" + descriptor, path = tempfile.mkstemp(prefix="localai-media-", suffix=suffix) + try: + with os.fdopen(descriptor, "wb") as output: + descriptor = None + output.write(base64.b64decode(data)) + yield path + finally: + if descriptor is not None: + os.close(descriptor) + try: + os.remove(path) + except OSError: + pass + + +@contextlib.contextmanager +def cleanup_paths(): + """Collect temporary paths and remove them on success or failure.""" + paths = [] + try: + yield paths + finally: + for path in paths: + try: + os.remove(path) + except OSError: + pass diff --git a/backend/python/common/temp_utils_test.py b/backend/python/common/temp_utils_test.py new file mode 100644 index 000000000..eb743064a --- /dev/null +++ b/backend/python/common/temp_utils_test.py @@ -0,0 +1,41 @@ +import os +import tempfile +import unittest +from unittest import mock + +from temp_utils import cleanup_paths, materialize_base64 + + +class MaterializeBase64Test(unittest.TestCase): + def test_removes_materialized_file_after_success(self): + with tempfile.TemporaryDirectory() as directory: + with mock.patch.object(tempfile, "tempdir", directory): + with materialize_base64("aGVsbG8=", suffix=".data") as path: + with open(path, "rb") as materialized: + self.assertEqual(materialized.read(), b"hello") + self.assertFalse(os.path.exists(path)) + + def test_removes_materialized_file_when_consumer_fails(self): + with tempfile.TemporaryDirectory() as directory: + with mock.patch.object(tempfile, "tempdir", directory): + with self.assertRaisesRegex(RuntimeError, "decode failed"): + with materialize_base64("aGVsbG8="): + raise RuntimeError("decode failed") + self.assertEqual(os.listdir(directory), []) + + +class CleanupPathsTest(unittest.TestCase): + def test_removes_every_registered_path_after_failure(self): + with tempfile.TemporaryDirectory() as directory: + paths = [os.path.join(directory, name) for name in ("one.wav", "two.wav")] + with self.assertRaisesRegex(RuntimeError, "merge failed"): + with cleanup_paths() as registered: + for path in paths: + open(path, "wb").close() + registered.append(path) + raise RuntimeError("merge failed") + self.assertEqual(os.listdir(directory), []) + + +if __name__ == "__main__": + unittest.main() diff --git a/backend/python/longcat-video/backend.py b/backend/python/longcat-video/backend.py index 0b54dd71f..761381efe 100755 --- a/backend/python/longcat-video/backend.py +++ b/backend/python/longcat-video/backend.py @@ -6,6 +6,7 @@ import datetime import gc import math import os +import shutil import signal import subprocess import sys @@ -888,6 +889,13 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): def _release_model(self): self.pipeline = None self.model_kind = None + try: + if hasattr(self, "dist") and self.dist.is_initialized(): + self.dist.destroy_process_group() + finally: + if self._dist_store_dir is not None: + shutil.rmtree(self._dist_store_dir, ignore_errors=True) + self._dist_store_dir = None gc.collect() if hasattr(self, "torch") and self.torch.cuda.is_available(): self.torch.cuda.empty_cache() diff --git a/backend/python/vllm-omni/backend.py b/backend/python/vllm-omni/backend.py index cc426e144..a23a1e535 100644 --- a/backend/python/vllm-omni/backend.py +++ b/backend/python/vllm-omni/backend.py @@ -19,7 +19,6 @@ import base64 import io import json import gc -import tempfile from PIL import Image import torch @@ -34,6 +33,7 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'common')) sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'common')) from grpc_auth import get_auth_interceptors from model_utils import resolve_model_reference +from temp_utils import materialize_base64 from vllm_utils import parse_options, messages_to_dicts, setup_parsers @@ -118,13 +118,8 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): return video_to_ndarrays(video_path, num_frames=16) # Try base64 decode try: - timestamp = str(int(time.time() * 1000)) - p = os.path.join(tempfile.gettempdir(), f"vl-{timestamp}.data") - with open(p, "wb") as f: - f.write(base64.b64decode(video_path)) - video = VideoAsset(name=p).np_ndarrays - os.remove(p) - return video + with materialize_base64(video_path, suffix=".data") as path: + return VideoAsset(name=path).np_ndarrays except: return None @@ -136,15 +131,9 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): return (audio_signal.astype(np.float32), sr) # Try base64 decode try: - audio_data = base64.b64decode(audio_path) - # Save to temp file and load - timestamp = str(int(time.time() * 1000)) - p = os.path.join(tempfile.gettempdir(), f"audio-{timestamp}.wav") - with open(p, "wb") as f: - f.write(audio_data) - audio_signal, sr = librosa.load(p, sr=16000) - os.remove(p) - return (audio_signal.astype(np.float32), sr) + with materialize_base64(audio_path, suffix=".wav") as path: + audio_signal, sr = librosa.load(path, sr=16000) + return (audio_signal.astype(np.float32), sr) except: return None diff --git a/backend/python/vllm/backend.py b/backend/python/vllm/backend.py index de9af3798..6dabbf4bf 100644 --- a/backend/python/vllm/backend.py +++ b/backend/python/vllm/backend.py @@ -10,7 +10,6 @@ import os import json import time import gc -import tempfile from typing import List from PIL import Image @@ -23,6 +22,7 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'common')) from python_utils import attach_media_parts from grpc_auth import get_auth_interceptors from model_utils import resolve_model_reference +from temp_utils import materialize_base64 from vllm_utils import apply_options_to_engine_args, normalize_option_key from vllm.engine.arg_utils import AsyncEngineArgs @@ -1005,13 +1005,8 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): Video: The loaded video. """ try: - timestamp = str(int(time.time() * 1000)) # Generate timestamp - p = os.path.join(tempfile.gettempdir(), f"vl-{timestamp}.data") - with open(p, "wb") as f: - f.write(base64.b64decode(video_path)) - video = VideoAsset(name=p).np_ndarrays - os.remove(p) - return video + with materialize_base64(video_path, suffix=".data") as path: + return VideoAsset(name=path).np_ndarrays except Exception as e: print(f"Error loading video {video_path}: {e}", file=sys.stderr) return None diff --git a/core/services/worker/supervisor.go b/core/services/worker/supervisor.go index 184419457..1d9b52365 100644 --- a/core/services/worker/supervisor.go +++ b/core/services/worker/supervisor.go @@ -597,6 +597,7 @@ func (s *backendSupervisor) reapDeadProcess(key string, bp *backendProcess) { if bp == nil { return } + s.cleanupProcessRuntime(bp.proc) if bp.port <= 0 { xlog.Error("Cannot recycle backend port: dead process has invalid recorded port", "backend", key, "addr", bp.addr, "port", bp.port) return @@ -614,6 +615,7 @@ func (s *backendSupervisor) releaseBackendStart(key string, bp *backendProcess) return } delete(s.processes, key) + s.cleanupProcessRuntime(bp.proc) if bp.port <= 0 { xlog.Error("Cannot recycle backend port: startup has invalid recorded port", "backend", key, "addr", bp.addr, "port", bp.port) return @@ -947,6 +949,7 @@ func (s *backendSupervisor) finishBackendStop(key string, bp *backendProcess, st return fmt.Errorf("stopping backend process %s: %w", key, stopErr) } delete(s.processes, key) + s.cleanupProcessRuntime(bp.proc) if bp.port <= 0 { xlog.Error("Cannot recycle backend port: process has invalid recorded port", "backend", key, "addr", bp.addr, "port", bp.port) return nil @@ -955,6 +958,14 @@ func (s *backendSupervisor) finishBackendStop(key string, bp *backendProcess, st return nil } +func (s *backendSupervisor) cleanupProcessRuntime(proc *process.Process) { + // Some focused supervisor tests provide synthetic process handles without a + // ModelLoader. Production processes always come from s.ml.StartProcess. + if s.ml != nil { + s.ml.CleanupProcessRuntime(proc) + } +} + // stopAllBackends stops all running backend processes. func (s *backendSupervisor) stopAllBackends(force bool) { s.mu.Lock() diff --git a/docs/content/reference/cli-reference.md b/docs/content/reference/cli-reference.md index 2ba6c96e8..d00d8046b 100644 --- a/docs/content/reference/cli-reference.md +++ b/docs/content/reference/cli-reference.md @@ -27,9 +27,16 @@ Complete reference for all LocalAI command-line interface (CLI) parameters and e | `--upload-path` | `TMPDIR/localai-UID/upload` | Path to store uploads from files API. Defaults under the OS temp dir (`$TMPDIR`, falling back to `/tmp`), scoped to the current user's UID. | `$LOCALAI_UPLOAD_PATH`, `$UPLOAD_PATH` | | `--localai-config-dir` | `BASEPATH/configuration` | Directory for dynamic loading of certain configuration files (currently runtime_settings.json, api_keys.json, and external_backends.json). See [Runtime Settings]({{%relref "features/runtime-settings" %}}) for web-based configuration. | `$LOCALAI_CONFIG_DIR` | | `--localai-config-dir-poll-interval` | | Time duration to poll the LocalAI Config Dir if your system has broken fsnotify events (example: `1m`) | `$LOCALAI_CONFIG_DIR_POLL_INTERVAL` | + | `--models-config-file` | | YAML file containing a list of model backend configs (alias: `--config-file`) | `$LOCALAI_MODELS_CONFIG_FILE`, `$CONFIG_FILE` | | `--artifact-download-concurrency` | `1` | How many files of a model artifact to download at once. `1` downloads sequentially. Raising it helps artifacts split into many files on a fast link, at the cost of more concurrent load on the models volume. Whole files only — a single file is never split, so resume and per-file checksum verification are unaffected | `$LOCALAI_ARTIFACT_DOWNLOAD_CONCURRENCY` | +Backend processes receive a private scratch directory through `TMPDIR`, `TMP`, +and `TEMP`. LocalAI removes that directory when the backend exits and removes +abandoned directories left by a LocalAI crash before starting another backend. +Set `$LOCALAI_BACKEND_TEMP_DIR` to choose their base volume. LocalAI always +appends `localai-UID/backend-runtime`; the default base is `TMPDIR`. + ## Backend Flags | Parameter | Default | Description | Environment Variable | diff --git a/pkg/model/initializers.go b/pkg/model/initializers.go index e73d45535..49276abc3 100644 --- a/pkg/model/initializers.go +++ b/pkg/model/initializers.go @@ -173,7 +173,7 @@ func (ml *ModelLoader) spawnGRPCModel(backend, uri string, o *Options, modelID, if !ready { xlog.Debug("GRPC Service NOT ready") startupErr := grpcStartupError(client.Process()) - stopLoadProcess(client, modelID) + ml.stopLoadProcess(client, modelID) return nil, startupErr } @@ -189,11 +189,11 @@ func (ml *ModelLoader) spawnGRPCModel(backend, uri string, o *Options, modelID, res, err := client.GRPC(o.parallelRequests, ml.wd).LoadModel(o.context, options) if err != nil { - stopLoadProcess(client, modelID) + ml.stopLoadProcess(client, modelID) return nil, fmt.Errorf("could not load model: %w", err) } if !res.Success { - stopLoadProcess(client, modelID) + ml.stopLoadProcess(client, modelID) return nil, fmt.Errorf("could not load model (no success): %s", res.Message) } @@ -260,7 +260,7 @@ func lastNonEmptyLine(path string, maxBytes int64) string { // stopLoadProcess tears down a backend process whose load did not complete. // The stop error is only logged: the load error is what the caller reports. -func stopLoadProcess(client *Model, modelID string) { +func (ml *ModelLoader) stopLoadProcess(client *Model, modelID string) { process := client.Process() if process == nil { return @@ -268,6 +268,7 @@ func stopLoadProcess(client *Model, modelID string) { if err := process.Stop(); err != nil { xlog.Warn("failed to stop backend process after failed load", "error", err, "modelID", modelID) } + ml.cleanupProcessRuntime(process) } // parallelSlotsFromOptions returns the effective n_parallel from the backend diff --git a/pkg/model/loader.go b/pkg/model/loader.go index 322b11e36..2df4aee2a 100644 --- a/pkg/model/loader.go +++ b/pkg/model/loader.go @@ -106,6 +106,10 @@ type ModelLoader struct { // the exit code can't, since a child killed by our own SIGTERM/SIGKILL // reports -1, indistinguishable from a signal-induced crash. stoppingProcs sync.Map + // processRuntimes keeps the owned state/scratch directory alive until the + // loader has consumed any exit diagnostics. The exit watcher removes the + // potentially large scratch contents immediately. + processRuntimes sync.Map // loadFailures records, per modelID, the cooldown window applied after a // failed load so that a client repeatedly polling a broken model does not // spawn (and leak) a fresh backend process on every request. Guarded by mu. diff --git a/pkg/model/process.go b/pkg/model/process.go index fb1ed004a..9b799e5a5 100644 --- a/pkg/model/process.go +++ b/pkg/model/process.go @@ -177,12 +177,14 @@ func (ml *ModelLoader) deleteProcess(ctx context.Context, s string, force bool) // A concurrently crashed/already-reaped process can no longer own // resources even if Stop could not read or signal its PID. store.Delete(s) + ml.cleanupProcessRuntime(process) return nil } return err } store.Delete(s) + ml.cleanupProcessRuntime(process) return nil } func (ml *ModelLoader) StopGRPC(filter GRPCProcessFilter) error { @@ -231,16 +233,6 @@ func (ml *ModelLoader) StartProcess(grpcProcess, id string, serverAddress string return ml.startProcess(grpcProcess, id, serverAddress, args...) } -// newProcessStateDir creates the directory a backend process uses for its pid, -// state and log files, and reports why when it cannot. -func newProcessStateDir() (string, error) { - dir, err := os.MkdirTemp(os.TempDir(), "go-processmanager") - if err != nil { - return "", fmt.Errorf("creating backend process state directory under %s: %w", os.TempDir(), err) - } - return dir, nil -} - func (ml *ModelLoader) startProcess(grpcProcess, id string, serverAddress string, args ...string) (*process.Process, error) { // Make sure the process is executable // Check first if it has executable permissions @@ -262,7 +254,12 @@ func (ml *ModelLoader) startProcess(grpcProcess, id string, serverAddress string return nil, err } - env := os.Environ() + runtime, err := newBackendProcessRuntime() + if err != nil { + return nil, err + } + + env := backendTempEnvironment(os.Environ(), runtime.tempDir) // Vulkan backends are self-contained: they bundle their own loader and // Mesa driver .so files in lib/ plus the matching ICD manifests in // vulkan/icd.d/. Point the loader at those manifests so it doesn't rely on @@ -271,16 +268,14 @@ func (ml *ModelLoader) startProcess(grpcProcess, id string, serverAddress string // and the GPU would silently fall back to CPU). No-op for other backends. env = append(env, vulkanICDEnv(workDir)...) - // Resolve the state directory here rather than through + // Resolve and own the state directory here rather than through // process.WithTemporaryStateDir(). process.New applies its options but // discards the error they return, so a temp directory that cannot be // created leaves StateDir empty and every later option unapplied. Run() - // then reported "mkdir : no such file or directory" with no path, hiding - // the real cause (a full volume, or a TMPDIR that no longer resolves). - stateDir, err := newProcessStateDir() - if err != nil { - return nil, err - } + // then reports "mkdir : no such file or directory" with no useful path. + // The same owned directory also contains backend scratch so an unexpected + // exit cannot strand request files directly in the host's shared /tmp. + stateDir := runtime.dir grpcControlProcess := process.New( process.WithStateDir(stateDir), @@ -296,8 +291,10 @@ func (ml *ModelLoader) startProcess(grpcProcess, id string, serverAddress string } if err := grpcControlProcess.Run(); err != nil { + runtime.cleanup() return grpcControlProcess, err } + ml.processRuntimes.Store(grpcControlProcess, runtime) xlog.Debug("GRPC Service state dir", "dir", grpcControlProcess.StateDir()) @@ -376,11 +373,35 @@ func (ml *ModelLoader) startProcess(grpcProcess, id string, serverAddress string } xlog.Warn("Backend process exited unexpectedly", fields...) } + runtime.cleanupScratch() + close(runtime.diagnosticsDone) }() return grpcControlProcess, nil } +func (ml *ModelLoader) cleanupProcessRuntime(process *process.Process) { + if process == nil { + return + } + value, ok := ml.processRuntimes.LoadAndDelete(process) + if !ok { + return + } + runtime := value.(*backendProcessRuntime) + go func() { + <-runtime.diagnosticsDone + runtime.cleanup() + }() +} + +// CleanupProcessRuntime releases state and scratch owned by a process started +// through StartProcess. Callers that supervise processes outside ModelLoader's +// model store must invoke it after they have consumed exit diagnostics. +func (ml *ModelLoader) CleanupProcessRuntime(process *process.Process) { + ml.cleanupProcessRuntime(process) +} + // vulkanICDEnv returns environment overrides that point the Vulkan loader at // the ICD manifests a backend bundles in /vulkan/icd.d. Vulkan // backends ship a self-contained stack — their own loader and Mesa driver .so diff --git a/pkg/model/process_exit_test.go b/pkg/model/process_exit_test.go index cc5554bcf..2b941bf6e 100644 --- a/pkg/model/process_exit_test.go +++ b/pkg/model/process_exit_test.go @@ -15,8 +15,10 @@ import ( var _ = Describe("backend process exit diagnostics", func() { It("includes the exit code and final stderr line for an unexpected exit", func() { tmpDir := GinkgoT().TempDir() + backendTempRoot := filepath.Join(tmpDir, "backend-runtime") + GinkgoT().Setenv(backendTempDirEnv, backendTempRoot) backendPath := filepath.Join(tmpDir, "failing-backend") - Expect(os.WriteFile(backendPath, []byte("#!/bin/sh\necho 'first diagnostic' >&2\necho 'fatal metal pipeline error' >&2\nexit 42\n"), 0o700)).To(Succeed()) + Expect(os.WriteFile(backendPath, []byte("#!/bin/sh\nprintf '%s' \"$TMPDIR\" > \"$0.tmpdir\"\necho 'first diagnostic' >&2\necho 'fatal metal pipeline error' >&2\nexit 42\n"), 0o700)).To(Succeed()) captured := &bytes.Buffer{} handler := slog.NewTextHandler(captured, &slog.HandlerOptions{Level: slog.LevelWarn}) @@ -29,10 +31,16 @@ var _ = Describe("backend process exit diagnostics", func() { process, err := loader.startProcess(backendPath, "test-model", "127.0.0.1:65535") Expect(err).ToNot(HaveOccurred()) Eventually(process.Done()).Should(BeClosed()) + backendTemp, err := os.ReadFile(backendPath + ".tmpdir") + Expect(err).ToNot(HaveOccurred()) + Expect(string(backendTemp)).To(Equal(filepath.Join(process.StateDir(), "tmp"))) + Eventually(string(backendTemp)).ShouldNot(BeADirectory()) Eventually(captured.String).Should(And( ContainSubstring("Backend process exited unexpectedly"), ContainSubstring("exitCode=42"), ContainSubstring(`stderr="fatal metal pipeline error"`), )) + loader.cleanupProcessRuntime(process) + Eventually(process.StateDir()).ShouldNot(BeADirectory()) }) }) diff --git a/pkg/model/process_runtime.go b/pkg/model/process_runtime.go new file mode 100644 index 000000000..b7443ef7d --- /dev/null +++ b/pkg/model/process_runtime.go @@ -0,0 +1,168 @@ +package model + +import ( + "fmt" + "os" + "path/filepath" + "strings" + "sync" + + "github.com/gofrs/flock" + "github.com/mudler/xlog" +) + +const ( + backendTempDirEnv = "LOCALAI_BACKEND_TEMP_DIR" + backendRuntimeDirPrefix = "process-" + backendRuntimeMarker = ".localai-backend-runtime" + backendRuntimeMagic = "localai-backend-runtime-v1\n" +) + +// backendProcessRuntime owns both go-processmanager's state and all temporary +// files created by one backend process. The held lock distinguishes a live +// runtime from one abandoned when LocalAI was killed or crashed. +type backendProcessRuntime struct { + dir string + tempDir string + lock *flock.Flock + scratch sync.Once + once sync.Once + // diagnosticsDone closes after the exit watcher has read the state files. + diagnosticsDone chan struct{} +} + +func backendRuntimeRoot() string { + base := os.TempDir() + if configured := os.Getenv(backendTempDirEnv); configured != "" { + base = configured + } + // Always append a LocalAI- and user-specific namespace. Even if an operator + // points the configurable base at /tmp, the sweeper never inspects unrelated + // process-* directories in that shared parent. + return filepath.Join(base, fmt.Sprintf("localai-%d", os.Getuid()), "backend-runtime") +} + +func newBackendProcessRuntime() (*backendProcessRuntime, error) { + root := backendRuntimeRoot() + if err := os.MkdirAll(root, 0o700); err != nil { + return nil, fmt.Errorf("creating backend runtime root %s: %w", root, err) + } + + // Serialize sweeping with creation. Otherwise a second LocalAI instance + // could observe the new directory in the tiny window before its owner lock + // is acquired and mistake it for an abandoned runtime. + sweepLock := flock.New(filepath.Join(root, ".sweep.lock")) + if err := sweepLock.Lock(); err != nil { + return nil, fmt.Errorf("locking backend runtime root %s: %w", root, err) + } + defer func() { + if err := sweepLock.Unlock(); err != nil { + xlog.Warn("Failed to unlock backend runtime root", "root", root, "error", err) + } + }() + + sweepAbandonedBackendRuntimes(root) + + dir, err := os.MkdirTemp(root, backendRuntimeDirPrefix) + if err != nil { + return nil, fmt.Errorf("creating backend process runtime under %s: %w", root, err) + } + if err := os.WriteFile(filepath.Join(dir, backendRuntimeMarker), []byte(backendRuntimeMagic), 0o600); err != nil { + _ = os.RemoveAll(dir) + return nil, fmt.Errorf("marking backend process runtime %s: %w", dir, err) + } + runtimeLock := flock.New(filepath.Join(dir, ".owner.lock")) + if err := runtimeLock.Lock(); err != nil { + _ = os.RemoveAll(dir) + return nil, fmt.Errorf("locking backend process runtime %s: %w", dir, err) + } + tempDir := filepath.Join(dir, "tmp") + if err := os.Mkdir(tempDir, 0o700); err != nil { + _ = runtimeLock.Unlock() + _ = os.RemoveAll(dir) + return nil, fmt.Errorf("creating backend scratch directory %s: %w", tempDir, err) + } + + return &backendProcessRuntime{ + dir: dir, + tempDir: tempDir, + lock: runtimeLock, + diagnosticsDone: make(chan struct{}), + }, nil +} + +func sweepAbandonedBackendRuntimes(root string) { + entries, err := os.ReadDir(root) + if err != nil { + xlog.Warn("Failed to inspect backend runtime root", "root", root, "error", err) + return + } + for _, entry := range entries { + if !entry.IsDir() || !strings.HasPrefix(entry.Name(), backendRuntimeDirPrefix) { + continue + } + dir := filepath.Join(root, entry.Name()) + marker, err := os.ReadFile(filepath.Join(dir, backendRuntimeMarker)) + if err != nil || string(marker) != backendRuntimeMagic { + continue + } + ownerLock := flock.New(filepath.Join(dir, ".owner.lock")) + available, err := ownerLock.TryLock() + if err != nil { + xlog.Warn("Failed to inspect backend runtime ownership", "dir", dir, "error", err) + continue + } + if !available { + continue + } + if err := ownerLock.Unlock(); err != nil { + xlog.Warn("Failed to release abandoned backend runtime lock", "dir", dir, "error", err) + continue + } + if err := os.RemoveAll(dir); err != nil { + xlog.Warn("Failed to remove abandoned backend runtime", "dir", dir, "error", err) + } + } +} + +func (r *backendProcessRuntime) cleanup() { + if r == nil { + return + } + r.once.Do(func() { + r.cleanupScratch() + if err := r.lock.Unlock(); err != nil { + xlog.Warn("Failed to unlock backend process runtime", "dir", r.dir, "error", err) + } + if err := os.RemoveAll(r.dir); err != nil { + xlog.Warn("Failed to remove backend process runtime", "dir", r.dir, "error", err) + } + }) +} + +func (r *backendProcessRuntime) cleanupScratch() { + if r == nil { + return + } + r.scratch.Do(func() { + if err := os.RemoveAll(r.tempDir); err != nil { + xlog.Warn("Failed to remove backend scratch directory", "dir", r.tempDir, "error", err) + } + }) +} + +func backendTempEnvironment(env []string, tempDir string) []string { + result := make([]string, 0, len(env)+3) + for _, entry := range env { + key, _, found := strings.Cut(entry, "=") + if found && (key == "TMPDIR" || key == "TMP" || key == "TEMP") { + continue + } + result = append(result, entry) + } + return append(result, + "TMPDIR="+tempDir, + "TMP="+tempDir, + "TEMP="+tempDir, + ) +} diff --git a/pkg/model/process_runtime_test.go b/pkg/model/process_runtime_test.go new file mode 100644 index 000000000..b27547300 --- /dev/null +++ b/pkg/model/process_runtime_test.go @@ -0,0 +1,109 @@ +package model + +import ( + "os" + "path/filepath" + "strings" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Backend process runtime directory", func() { + It("keeps active runtimes while sweeping abandoned ones", func() { + root := GinkgoT().TempDir() + GinkgoT().Setenv(backendTempDirEnv, root) + ownedRoot := backendRuntimeRoot() + unrelated := filepath.Join(root, backendRuntimeDirPrefix+"unrelated") + Expect(os.MkdirAll(unrelated, 0o700)).To(Succeed()) + Expect(os.WriteFile(filepath.Join(unrelated, "keep"), []byte("unrelated"), 0o600)).To(Succeed()) + + active, err := newBackendProcessRuntime() + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(active.cleanup) + Expect(os.WriteFile(filepath.Join(active.tempDir, "active.img"), []byte("active"), 0o600)).To(Succeed()) + + abandoned := filepath.Join(ownedRoot, backendRuntimeDirPrefix+"abandoned") + Expect(os.MkdirAll(abandoned, 0o700)).To(Succeed()) + Expect(os.WriteFile(filepath.Join(abandoned, backendRuntimeMarker), []byte(backendRuntimeMagic), 0o600)).To(Succeed()) + Expect(os.WriteFile(filepath.Join(abandoned, "orphan.img"), []byte("orphan"), 0o600)).To(Succeed()) + foreign := filepath.Join(ownedRoot, backendRuntimeDirPrefix+"foreign") + Expect(os.MkdirAll(foreign, 0o700)).To(Succeed()) + Expect(os.WriteFile(filepath.Join(foreign, "keep"), []byte("foreign"), 0o600)).To(Succeed()) + + other, err := newBackendProcessRuntime() + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(other.cleanup) + + Expect(active.dir).To(BeADirectory()) + Expect(abandoned).ToNot(BeAnExistingFile()) + Expect(foreign).To(BeADirectory()) + Expect(unrelated).To(BeADirectory()) + }) + + It("uses one private directory for process state and backend scratch", func() { + root := GinkgoT().TempDir() + GinkgoT().Setenv(backendTempDirEnv, root) + + runtime, err := newBackendProcessRuntime() + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(runtime.cleanup) + + Expect(filepath.Dir(runtime.dir)).To(Equal(backendRuntimeRoot())) + Expect(runtime.tempDir).To(Equal(filepath.Join(runtime.dir, "tmp"))) + info, err := os.Stat(runtime.tempDir) + Expect(err).ToNot(HaveOccurred()) + Expect(info.IsDir()).To(BeTrue()) + Expect(info.Mode().Perm()).To(Equal(os.FileMode(0o700))) + }) + + It("overrides inherited temp variables for the backend only", func() { + env := backendTempEnvironment([]string{ + "PATH=/bin", + "TMPDIR=/old/tmpdir", + "TMP=/old/tmp", + "TEMP=/old/temp", + }, "/owned/scratch") + + Expect(env).To(ConsistOf( + "PATH=/bin", + "TMPDIR=/owned/scratch", + "TMP=/owned/scratch", + "TEMP=/owned/scratch", + )) + for _, key := range []string{"TMPDIR", "TMP", "TEMP"} { + count := 0 + for _, entry := range env { + if strings.HasPrefix(entry, key+"=") { + count++ + } + } + Expect(count).To(Equal(1), key) + } + }) + + It("removes the runtime when its owner exits", func() { + GinkgoT().Setenv(backendTempDirEnv, GinkgoT().TempDir()) + + runtime, err := newBackendProcessRuntime() + Expect(err).ToNot(HaveOccurred()) + dir := runtime.dir + runtime.cleanup() + + Expect(dir).ToNot(BeAnExistingFile()) + }) + + It("reports which configured root cannot be used", func() { + parent := GinkgoT().TempDir() + file := filepath.Join(parent, "not-a-directory") + Expect(os.WriteFile(file, []byte("x"), 0o600)).To(Succeed()) + base := filepath.Join(file, "backend-runtime") + GinkgoT().Setenv(backendTempDirEnv, base) + root := backendRuntimeRoot() + + runtime, err := newBackendProcessRuntime() + Expect(err).To(HaveOccurred()) + Expect(runtime).To(BeNil()) + Expect(err.Error()).To(ContainSubstring(root)) + }) +}) diff --git a/pkg/model/process_statedir_test.go b/pkg/model/process_statedir_test.go deleted file mode 100644 index 207c2991d..000000000 --- a/pkg/model/process_statedir_test.go +++ /dev/null @@ -1,38 +0,0 @@ -package model - -import ( - "os" - "path/filepath" - - . "github.com/onsi/ginkgo/v2" - . "github.com/onsi/gomega" -) - -var _ = Describe("Backend process state directory", func() { - It("reports why the state directory could not be created", func() { - // A worker whose volume is full, or whose TMPDIR no longer resolves, - // cannot get a state directory. go-processmanager's New() drops the - // option error, leaving StateDir empty, and Run() then failed with - // "mkdir : no such file or directory" naming no path at all. Resolving - // the directory here keeps the real cause attached. - GinkgoT().Setenv("TMPDIR", filepath.Join(GinkgoT().TempDir(), "does-not-exist")) - - dir, err := newProcessStateDir() - Expect(err).To(HaveOccurred()) - Expect(dir).To(BeEmpty()) - Expect(err.Error()).To(ContainSubstring("backend process state directory")) - Expect(err.Error()).To(ContainSubstring("does-not-exist"), - "the error must name the directory it could not create") - }) - - It("returns a usable directory when the temp location works", func() { - GinkgoT().Setenv("TMPDIR", GinkgoT().TempDir()) - - dir, err := newProcessStateDir() - Expect(err).ToNot(HaveOccurred()) - Expect(dir).ToNot(BeEmpty()) - info, statErr := os.Stat(dir) - Expect(statErr).ToNot(HaveOccurred()) - Expect(info.IsDir()).To(BeTrue()) - }) -})