Robert Wojciechowski
87c0cda9f6
Retain durable reservations until provider termination or explicit recovery, protect unresolved pools and credentials, and add scoped reconciliation. Consolidate lifecycle and secret policies and split prompt/media stages with regression coverage.
442 lines
16 KiB
Python
442 lines
16 KiB
Python
"""Validation and private storage for Runpod serverless worker settings.
|
|
|
|
Secret storage is supported only on POSIX masters. The module remains
|
|
importable elsewhere so ordinary worker configuration and redaction continue
|
|
to work, but secret-store operations fail explicitly without POSIX locking.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import copy
|
|
import json
|
|
import os
|
|
import re
|
|
import secrets
|
|
import stat
|
|
import threading
|
|
from contextlib import contextmanager
|
|
from pathlib import Path
|
|
from typing import Iterator
|
|
|
|
from .secret_fields import is_secret_field_name
|
|
|
|
try:
|
|
import fcntl
|
|
except ImportError: # Windows and other platforms without POSIX file locking.
|
|
fcntl = None
|
|
|
|
|
|
class ServerlessConfigError(ValueError):
|
|
"""A serverless worker configuration is invalid."""
|
|
|
|
|
|
class SecretStoreError(RuntimeError):
|
|
"""The secret store could not be accessed safely."""
|
|
|
|
|
|
_ENDPOINT_ID_RE = re.compile(r"[A-Za-z0-9_-]{1,200}", re.ASCII)
|
|
_UNSUPPORTED_SECRET_STORE_MESSAGE = (
|
|
"Serverless credential storage is currently supported on POSIX masters only"
|
|
)
|
|
_NUMERIC_FIELDS = {
|
|
"jobs_per_run": (1, 16, 1),
|
|
"concurrency": (1, 16, 1),
|
|
"timeout_seconds": (30, 86400, 1800),
|
|
}
|
|
_THREAD_LOCKS: dict[str, threading.Lock] = {}
|
|
_THREAD_LOCKS_GUARD = threading.Lock()
|
|
|
|
|
|
def validate_serverless_worker(worker: dict) -> dict:
|
|
"""Return a validated copy of a Runpod serverless worker record.
|
|
|
|
Provider credentials are intentionally not required here; they are stored
|
|
separately by :class:`SecretStore`.
|
|
"""
|
|
if not isinstance(worker, dict):
|
|
raise ServerlessConfigError("worker must be an object")
|
|
|
|
try:
|
|
result = copy.deepcopy(worker)
|
|
except Exception:
|
|
raise ServerlessConfigError("worker must be copyable") from None
|
|
|
|
worker_id = result.get("id")
|
|
if not isinstance(worker_id, str) or not worker_id.strip():
|
|
raise ServerlessConfigError("id must be a non-empty string")
|
|
|
|
endpoint_id = result.get("endpoint_id")
|
|
if not isinstance(endpoint_id, str) or not _ENDPOINT_ID_RE.fullmatch(endpoint_id):
|
|
raise ServerlessConfigError("endpoint_id must be an opaque Runpod identifier")
|
|
|
|
if result.get("type", "serverless") != "serverless":
|
|
raise ServerlessConfigError("type must be serverless")
|
|
if result.get("provider", "runpod") != "runpod":
|
|
raise ServerlessConfigError("provider is not supported")
|
|
if "enabled" in result and type(result["enabled"]) is not bool:
|
|
raise ServerlessConfigError("enabled must be a boolean")
|
|
|
|
for field, (minimum, maximum, default) in _NUMERIC_FIELDS.items():
|
|
value = result.get(field, default)
|
|
if type(value) is not int or not minimum <= value <= maximum:
|
|
raise ServerlessConfigError(f"{field} must be an integer from {minimum} to {maximum}")
|
|
result[field] = value
|
|
|
|
result["type"] = "serverless"
|
|
result["provider"] = "runpod"
|
|
return result
|
|
|
|
|
|
def _thread_lock_for(path: str) -> threading.Lock:
|
|
with _THREAD_LOCKS_GUARD:
|
|
lock = _THREAD_LOCKS.get(path)
|
|
if lock is None:
|
|
lock = threading.Lock()
|
|
_THREAD_LOCKS[path] = lock
|
|
return lock
|
|
|
|
|
|
class SecretStore:
|
|
"""Atomically persist worker-id to API-key mappings in a private JSON file."""
|
|
|
|
def __init__(self, path: Path):
|
|
self.path = Path(path)
|
|
|
|
def get(self, worker_id: str) -> str | None:
|
|
self._require_supported()
|
|
worker_id = self._validate_worker_id(worker_id)
|
|
try:
|
|
with self._locked(create_parent=False) as (parent_fd, filename):
|
|
data = self._read_mapping(parent_fd, filename)
|
|
return data.get(worker_id)
|
|
except FileNotFoundError:
|
|
return None
|
|
except SecretStoreError:
|
|
raise
|
|
except Exception:
|
|
raise SecretStoreError("Unable to access secret store") from None
|
|
|
|
def set(self, worker_id: str, key: str) -> None:
|
|
self._require_supported()
|
|
worker_id = self._validate_worker_id(worker_id)
|
|
if (
|
|
not isinstance(key, str)
|
|
or not key
|
|
or len(key) > 4096
|
|
or "\r" in key
|
|
or "\n" in key
|
|
):
|
|
raise SecretStoreError("secret must be 1 to 4096 characters and contain no CR or LF")
|
|
try:
|
|
with self._locked(create_parent=True) as (parent_fd, filename):
|
|
data = self._read_mapping(parent_fd, filename)
|
|
data[worker_id] = key
|
|
self._write_mapping(parent_fd, filename, data)
|
|
except SecretStoreError:
|
|
raise
|
|
except Exception:
|
|
raise SecretStoreError("Unable to access secret store") from None
|
|
|
|
def delete(self, worker_id: str) -> None:
|
|
self._require_supported()
|
|
worker_id = self._validate_worker_id(worker_id)
|
|
try:
|
|
with self._locked(create_parent=False) as (parent_fd, filename):
|
|
data = self._read_mapping(parent_fd, filename)
|
|
if worker_id in data:
|
|
del data[worker_id]
|
|
self._write_mapping(parent_fd, filename, data)
|
|
except FileNotFoundError:
|
|
return
|
|
except SecretStoreError:
|
|
raise
|
|
except Exception:
|
|
raise SecretStoreError("Unable to access secret store") from None
|
|
|
|
def has(self, worker_id: str) -> bool:
|
|
return self.get(worker_id) is not None
|
|
|
|
@staticmethod
|
|
def _require_supported() -> None:
|
|
# Never substitute a weaker Windows store: all operations retain the
|
|
# dir_fd/O_NOFOLLOW and flock protections implemented below.
|
|
if os.name != "posix" or fcntl is None:
|
|
raise SecretStoreError(_UNSUPPORTED_SECRET_STORE_MESSAGE)
|
|
|
|
@staticmethod
|
|
def _validate_worker_id(worker_id: str) -> str:
|
|
if not isinstance(worker_id, str) or not worker_id.strip():
|
|
raise SecretStoreError("worker_id must be a non-empty string")
|
|
return worker_id
|
|
|
|
def _absolute_path(self) -> Path:
|
|
path = self.path
|
|
if ".." in path.parts:
|
|
raise SecretStoreError("unsafe secret store path")
|
|
if not path.is_absolute():
|
|
path = Path.cwd() / path
|
|
if len(path.parts) < 2 or path.name in ("", ".", ".."):
|
|
raise SecretStoreError("unsafe secret store path")
|
|
return path
|
|
|
|
@contextmanager
|
|
def _locked(self, create_parent: bool) -> Iterator[tuple[int, str]]:
|
|
lock_api = fcntl
|
|
if os.name != "posix" or lock_api is None:
|
|
raise SecretStoreError(_UNSUPPORTED_SECRET_STORE_MESSAGE)
|
|
path = self._absolute_path()
|
|
filename = path.name
|
|
lock_name = f".{filename}.lock"
|
|
thread_lock = _thread_lock_for(str(path))
|
|
|
|
with thread_lock:
|
|
parent_fd = self._open_parent(path, create=create_parent)
|
|
lock_fd = None
|
|
try:
|
|
flags = os.O_RDWR | os.O_CREAT | os.O_CLOEXEC | os.O_NOFOLLOW
|
|
lock_fd = os.open(lock_name, flags, 0o600, dir_fd=parent_fd)
|
|
lock_stat = os.fstat(lock_fd)
|
|
if not stat.S_ISREG(lock_stat.st_mode):
|
|
raise SecretStoreError("unsafe secret store lock")
|
|
os.fchmod(lock_fd, 0o600)
|
|
lock_api.flock(lock_fd, lock_api.LOCK_EX)
|
|
self._check_target(parent_fd, filename)
|
|
yield parent_fd, filename
|
|
except SecretStoreError:
|
|
raise
|
|
except Exception:
|
|
raise SecretStoreError("Unable to access secret store") from None
|
|
finally:
|
|
if lock_fd is not None:
|
|
try:
|
|
lock_api.flock(lock_fd, lock_api.LOCK_UN)
|
|
except OSError:
|
|
pass
|
|
os.close(lock_fd)
|
|
os.close(parent_fd)
|
|
|
|
@staticmethod
|
|
def _open_parent(path: Path, create: bool) -> int:
|
|
parts = path.parts
|
|
current_fd = os.open(path.anchor, os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC)
|
|
try:
|
|
for part in parts[1:-1]:
|
|
created = False
|
|
try:
|
|
next_fd = os.open(
|
|
part,
|
|
os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC | os.O_NOFOLLOW,
|
|
dir_fd=current_fd,
|
|
)
|
|
except FileNotFoundError:
|
|
if not create:
|
|
raise
|
|
try:
|
|
os.mkdir(part, 0o700, dir_fd=current_fd)
|
|
created = True
|
|
except FileExistsError:
|
|
pass
|
|
next_fd = os.open(
|
|
part,
|
|
os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC | os.O_NOFOLLOW,
|
|
dir_fd=current_fd,
|
|
)
|
|
if created:
|
|
os.fchmod(next_fd, 0o700)
|
|
os.close(current_fd)
|
|
current_fd = next_fd
|
|
return current_fd
|
|
except Exception:
|
|
os.close(current_fd)
|
|
raise
|
|
|
|
@staticmethod
|
|
def _check_target(parent_fd: int, filename: str) -> None:
|
|
try:
|
|
target_stat = os.stat(filename, dir_fd=parent_fd, follow_symlinks=False)
|
|
except FileNotFoundError:
|
|
return
|
|
if not stat.S_ISREG(target_stat.st_mode):
|
|
raise SecretStoreError("unsafe secret store file")
|
|
|
|
@classmethod
|
|
def _read_mapping(cls, parent_fd: int, filename: str) -> dict[str, str]:
|
|
cls._check_target(parent_fd, filename)
|
|
try:
|
|
fd = os.open(
|
|
filename,
|
|
os.O_RDONLY | os.O_CLOEXEC | os.O_NOFOLLOW,
|
|
dir_fd=parent_fd,
|
|
)
|
|
except FileNotFoundError:
|
|
return {}
|
|
except Exception:
|
|
raise SecretStoreError("Unable to access secret store") from None
|
|
|
|
try:
|
|
file_stat = os.fstat(fd)
|
|
if not stat.S_ISREG(file_stat.st_mode):
|
|
raise SecretStoreError("unsafe secret store file")
|
|
os.fchmod(fd, 0o600)
|
|
with os.fdopen(fd, "rb") as stream:
|
|
raw = stream.read()
|
|
try:
|
|
mapping = json.loads(raw.decode("utf-8"), object_pairs_hook=cls._unique_object)
|
|
except Exception:
|
|
raise SecretStoreError("secret store is corrupt") from None
|
|
if not isinstance(mapping, dict) or any(
|
|
not isinstance(worker_id, str)
|
|
or not worker_id.strip()
|
|
or not isinstance(key, str)
|
|
or not key
|
|
for worker_id, key in mapping.items()
|
|
):
|
|
raise SecretStoreError("secret store is corrupt")
|
|
return mapping
|
|
except SecretStoreError:
|
|
raise
|
|
except Exception:
|
|
raise SecretStoreError("Unable to access secret store") from None
|
|
finally:
|
|
try:
|
|
os.close(fd)
|
|
except OSError:
|
|
pass
|
|
|
|
@staticmethod
|
|
def _unique_object(pairs: list[tuple[str, object]]) -> dict[str, object]:
|
|
result: dict[str, object] = {}
|
|
for key, value in pairs:
|
|
if key in result:
|
|
raise ValueError("duplicate key")
|
|
result[key] = value
|
|
return result
|
|
|
|
@staticmethod
|
|
def _write_mapping(parent_fd: int, filename: str, mapping: dict[str, str]) -> None:
|
|
payload = json.dumps(mapping, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode("utf-8")
|
|
temp_name = f".{filename}.{secrets.token_hex(12)}.tmp"
|
|
temp_fd = None
|
|
try:
|
|
flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_CLOEXEC | os.O_NOFOLLOW
|
|
temp_fd = os.open(temp_name, flags, 0o600, dir_fd=parent_fd)
|
|
os.fchmod(temp_fd, 0o600)
|
|
with os.fdopen(temp_fd, "wb") as stream:
|
|
temp_fd = None
|
|
stream.write(payload)
|
|
stream.flush()
|
|
os.fsync(stream.fileno())
|
|
SecretStore._check_target(parent_fd, filename)
|
|
os.replace(temp_name, filename, src_dir_fd=parent_fd, dst_dir_fd=parent_fd)
|
|
os.fsync(parent_fd)
|
|
except SecretStoreError:
|
|
raise
|
|
except Exception:
|
|
raise SecretStoreError("Unable to access secret store") from None
|
|
finally:
|
|
if temp_fd is not None:
|
|
try:
|
|
os.close(temp_fd)
|
|
except OSError:
|
|
pass
|
|
try:
|
|
os.unlink(temp_name, dir_fd=parent_fd)
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
class ServerlessSecretChanges:
|
|
"""Apply serverless secret changes and restore them unless committed.
|
|
|
|
Use this around the associated config transaction so a failed config save
|
|
compensates any secret-store writes or deletions made along the way.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self._originals: dict[tuple[int, str], tuple[SecretStore, str, str | None]] = {}
|
|
self._mutations: list[tuple[int, str]] = []
|
|
self._committed = False
|
|
|
|
def __enter__(self) -> "ServerlessSecretChanges":
|
|
return self
|
|
|
|
def __exit__(self, _exc_type, _exc_value, _traceback) -> bool:
|
|
if not self._committed:
|
|
self.rollback()
|
|
return False
|
|
|
|
def capture(self, store: SecretStore, worker_id: str) -> str | None:
|
|
"""Capture and return the original key once per store/worker pair."""
|
|
mutation_id = (id(store), worker_id)
|
|
if mutation_id not in self._originals:
|
|
self._originals[mutation_id] = (store, worker_id, store.get(worker_id))
|
|
return self._originals[mutation_id][2]
|
|
|
|
def set(self, store: SecretStore, worker_id: str, key: str) -> None:
|
|
self._prepare_mutation(store, worker_id)
|
|
store.set(worker_id, key)
|
|
|
|
def delete(self, store: SecretStore, worker_id: str) -> None:
|
|
self._prepare_mutation(store, worker_id)
|
|
store.delete(worker_id)
|
|
|
|
def commit(self) -> None:
|
|
"""Keep the secret changes after the corresponding config save."""
|
|
self._committed = True
|
|
|
|
def rollback(self) -> None:
|
|
"""Best-effort restoration without masking the original route error."""
|
|
for mutation_id in reversed(self._mutations):
|
|
store, worker_id, prior_secret = self._originals[mutation_id]
|
|
try:
|
|
if prior_secret is None:
|
|
store.delete(worker_id)
|
|
else:
|
|
store.set(worker_id, prior_secret)
|
|
except Exception:
|
|
pass
|
|
self._mutations.clear()
|
|
|
|
def _prepare_mutation(self, store: SecretStore, worker_id: str) -> None:
|
|
self.capture(store, worker_id)
|
|
mutation_id = (id(store), worker_id)
|
|
if mutation_id not in self._mutations:
|
|
# Record before applying: a store operation may fail after replacing
|
|
# its file, so the exception path must still attempt compensation.
|
|
self._mutations.append(mutation_id)
|
|
|
|
|
|
def _redact_credentials(value):
|
|
if isinstance(value, dict):
|
|
redacted = {}
|
|
for key, item in value.items():
|
|
if is_secret_field_name(key):
|
|
continue
|
|
redacted[key] = _redact_credentials(item)
|
|
return redacted
|
|
if isinstance(value, list):
|
|
return [_redact_credentials(item) for item in value]
|
|
if isinstance(value, tuple):
|
|
return tuple(_redact_credentials(item) for item in value)
|
|
return value
|
|
|
|
|
|
def public_config(config, store: SecretStore):
|
|
"""Return a deep-copied public configuration with credentials removed."""
|
|
result = _redact_credentials(copy.deepcopy(config))
|
|
if not isinstance(result, dict):
|
|
return result
|
|
|
|
workers = result.get("workers")
|
|
if isinstance(workers, list):
|
|
for worker in workers:
|
|
if isinstance(worker, dict) and worker.get("type") == "serverless":
|
|
worker_id = worker.get("id")
|
|
configured = (
|
|
store.has(worker_id)
|
|
if isinstance(worker_id, str) and worker_id.strip()
|
|
else False
|
|
)
|
|
worker["api_key_configured"] = configured
|
|
return result
|