feat(BA-6155): add Valkey-backed chunk-store upload session engine

Introduce the concurrency-safe upload session engine that replaces the
single-file-append TUS model. Session metadata is the source of truth in
Valkey, guarded by a per-session lock (SET NX + token-compare Lua release);
only the chunk payload bytes live on the shared filesystem as per-offset
chunk_<offset>.dat files. Coordination across multiple Storage Proxy replicas
happens entirely through Valkey, so there is no dependency on filesystem lock
semantics (fcntl.flock) or NFS attribute-cache coherence; chunk payloads are
content-addressed by (offset, sha256) and idempotent, needing no coordination.

Contents:
- common/clients/valkey_client/valkey_tus: ValkeyTusClient — per-session
  state get/set (with TTL) + a per-session lock; reuses the Glide-based
  AbstractValkeyClient like the other valkey clients.
- common/defs: REDIS_TUS_DB; common/metrics: VALKEY_TUS layer.
- errors/upload.py: ChunkConflictError (409), UploadSessionCorruptedError (500)
- services/upload/tus_session.py:
  - ChunkRecord / SessionState (BackendAISchema; committed_offset as the
    largest contiguous prefix, missing_ranges, progress_percent) / ChunkAcceptance
  - UploadStatus StrEnum
  - TusUploadSession (ensure_initialized, read_state, write_temp_chunk,
    commit_chunk — idempotent dup / 409 conflict / no-op when completed,
    assemble, cleanup) taking a TusUploadSessionArgs that carries the
    ValkeyTusClient.
- Stale sessions auto-expire via the Valkey TTL (no separate GC sweep).
- Storage RootContext now provisions a ValkeyTusClient (server.py bootstrap).
- Integration tests drive a real Valkey (redis container) + tmp_path chunks.

Resolves BA-6154. Resolves BA-6155.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Author: jopemachine

changes/11766.feature.md

@@ -0,0 +1 @@
+Introduce `TusUploadSession`, a concurrency-safe upload engine for resumable (TUS) uploads that keeps session metadata and a per-session lock in Valkey (via `ValkeyTusClient`) while storing chunk payloads on the shared filesystem, so multiple storage-proxy replicas can accept chunks without corruption and without relying on filesystem lock semantics.

src/ai/backend/common/clients/valkey_client/valkey_tus/__init__.py

@@ -0,0 +1,5 @@
+"""Valkey client for TUS resumable-upload session metadata."""
+
+from .client import ValkeyTusClient
+
+__all__ = ["ValkeyTusClient"]

src/ai/backend/common/clients/valkey_client/valkey_tus/client.py

@@ -0,0 +1,111 @@
+"""
+Valkey client for TUS resumable-upload sessions.
+
+Stores per-session upload metadata (the source of truth) on the Glide-based
+valkey connection so that multiple Storage Proxy replicas share one view of an
+upload's progress without relying on shared-filesystem semantics. This client is
+metadata-only; the per-session distributed lock that serializes the
+read-modify-write window is a separate :class:`DistributedLockFactory` resource
+(see :mod:`ai.backend.storage.services.upload.tus_session`). The payload bytes
+stay on the shared filesystem; only this small metadata lives in Valkey.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Final, Self
+
+from glide import ExpirySet, ExpiryType
+
+from ai.backend.common.clients.valkey_client.client import (
+    AbstractValkeyClient,
+    create_valkey_client,
+)
+from ai.backend.common.exception import BackendAIError
+from ai.backend.common.metrics.metric import DomainType, LayerType
+from ai.backend.common.resilience import (
+    BackoffStrategy,
+    MetricArgs,
+    MetricPolicy,
+    Resilience,
+    RetryArgs,
+    RetryPolicy,
+)
+from ai.backend.common.types import TusSessionId, ValkeyTarget
+from ai.backend.logging import BraceStyleAdapter
+
+log = BraceStyleAdapter(logging.getLogger(__spec__.name))
+
+valkey_tus_resilience = Resilience(
+    policies=[
+        MetricPolicy(MetricArgs(domain=DomainType.VALKEY, layer=LayerType.VALKEY_TUS)),
+        RetryPolicy(
+            RetryArgs(
+                max_retries=3,
+                retry_delay=0.1,
+                backoff_strategy=BackoffStrategy.FIXED,
+                non_retryable_exceptions=(BackendAIError,),
+            )
+        ),
+    ]
+)
+
+_STATE_KEY_PREFIX: Final = "tus.upload.session"  # tus.upload.session:{session_id}
+
+# Stale (never-completed / abandoned) sessions are reclaimed by this TTL, so no
+# separate GC sweep is needed.
+_DEFAULT_STATE_TTL_SECONDS: Final = 24 * 60 * 60
+
+
+class ValkeyTusClient:
+    """Valkey-backed metadata store for TUS uploads."""
+
+    _client: AbstractValkeyClient
+
+    def __init__(self, client: AbstractValkeyClient) -> None:
+        self._client = client
+
+    @classmethod
+    async def create(
+        cls,
+        valkey_target: ValkeyTarget,
+        *,
+        db_id: int,
+        human_readable_name: str,
+    ) -> Self:
+        client = create_valkey_client(
+            valkey_target=valkey_target,
+            db_id=db_id,
+            human_readable_name=human_readable_name,
+        )
+        await client.connect()
+        return cls(client=client)
+
+    @valkey_tus_resilience.apply()
+    async def close(self) -> None:
+        await self._client.disconnect()
+
+    @staticmethod
+    def _state_key(session_id: TusSessionId) -> str:
+        return f"{_STATE_KEY_PREFIX}:{session_id}"
+
+    @valkey_tus_resilience.apply()
+    async def get_session_state(self, session_id: TusSessionId) -> bytes | None:
+        """Return the raw serialized session state, or ``None`` if absent."""
+        async with self._client.client() as conn:
+            return await conn.get(self._state_key(session_id))
+
+    @valkey_tus_resilience.apply()
+    async def set_session_state(
+        self,
+        session_id: TusSessionId,
+        payload: str | bytes,
+        *,
+        ttl_seconds: int = _DEFAULT_STATE_TTL_SECONDS,
+    ) -> None:
+        async with self._client.client() as conn:
+            await conn.set(
+                self._state_key(session_id),
+                payload,
+                expiry=ExpirySet(ExpiryType.SEC, ttl_seconds),
+            )

src/ai/backend/common/defs/__init__.py

@@ -11,6 +11,7 @@
 REDIS_STREAM_LOCK: Final = 5
 REDIS_CONTAINER_LOG: Final = 6
 REDIS_BGTASK_DB: Final = 7
+REDIS_TUS_DB: Final = 8
 
 
 class RedisRole(StrEnum):

src/ai/backend/common/lock.py

@@ -7,7 +7,7 @@
 from collections.abc import Mapping
 from io import IOBase
 from pathlib import Path
-from typing import Any, ClassVar
+from typing import Any, ClassVar, Protocol, runtime_checkable
 
 import trafaret as t
 from etcd_client import Client as EtcdClient
@@ -52,6 +52,11 @@ async def __aexit__(self, *exc_info: Any) -> bool | None:
         raise NotImplementedError
 
 
+@runtime_checkable
+class DistributedLockFactory(Protocol):
+    def __call__(self, lock_id: str, lifetime_hint: float) -> AbstractDistributedLock: ...
+
+
 class FileLock(AbstractDistributedLock):
     default_timeout: float = 3  # not allow infinite timeout for safety
 

src/ai/backend/common/metrics/metric.py

@@ -489,6 +489,7 @@ class LayerType(enum.StrEnum):
     VALKEY_STREAM = "valkey_stream"
     VALKEY_BGTASK = "valkey_bgtask"
     VALKEY_VOLUME_STATS = "valkey_volume_stats"
+    VALKEY_TUS = "valkey_tus"
 
     # Client layers
     AGENT_CLIENT = "agent_client"

src/ai/backend/common/types.py

@@ -395,6 +395,7 @@ def check_typed_tuple(value: tuple[Any, ...], types: tuple[type, ...]) -> tuple[
 ContainerPID = NewType("ContainerPID", PID)
 
 ContainerId = NewType("ContainerId", str)
+TusSessionId = NewType("TusSessionId", str)
 RuleId = NewType("RuleId", UUID)
 SessionId = NewType("SessionId", UUID)
 KernelId = NewType("KernelId", UUID)

src/ai/backend/storage/context.py

@@ -15,12 +15,14 @@
 from ai.backend.common.clients.valkey_client.valkey_artifact.client import (
     ValkeyArtifactDownloadTrackingClient,
 )
+from ai.backend.common.clients.valkey_client.valkey_tus import ValkeyTusClient
 from ai.backend.common.etcd import AsyncEtcd
 from ai.backend.common.events.dispatcher import (
     EventDispatcher,
     EventProducer,
 )
 from ai.backend.common.health_checker.probe import HealthProbe
+from ai.backend.common.lock import DistributedLockFactory
 from ai.backend.common.metrics.metric import CommonMetricRegistry
 from ai.backend.logging import BraceStyleAdapter
 
@@ -106,6 +108,8 @@ class RootContext:
     cors_options: Mapping[str, aiohttp_cors.ResourceOptions]
     manager_client_pool: ManagerHTTPClientPool
     valkey_artifact_client: ValkeyArtifactDownloadTrackingClient
+    valkey_tus_client: ValkeyTusClient
+    tus_lock_factory: DistributedLockFactory
     health_probe: HealthProbe
     volume_stats_observer: VolumeStatsObserver
     volume_stats_state: VolumeState

src/ai/backend/storage/errors/__init__.py

@@ -62,6 +62,10 @@
     QuotaScopeNotFoundError,
     QuotaTreeNotFoundError,
 )
+from .upload import (
+    ChunkConflictError,
+    UploadSessionCorruptedError,
+)
 from .vfolder import (
     InvalidSubpathError,
     VFolderNotFoundError,
@@ -90,6 +94,9 @@
     "InvalidDataLengthError",
     "ServiceNotInitializedError",
     "UploadOffsetMismatchError",
+    # upload
+    "ChunkConflictError",
+    "UploadSessionCorruptedError",
     # vfolder
     "VFolderNotFoundError",
     "InvalidSubpathError",

src/ai/backend/storage/errors/upload.py

@@ -0,0 +1,49 @@
+"""
+Upload session related exceptions.
+"""
+
+from __future__ import annotations
+
+from aiohttp import web
+
+from ai.backend.common.exception import (
+    BackendAIError,
+    ErrorCode,
+    ErrorDetail,
+    ErrorDomain,
+    ErrorOperation,
+)
+
+
+class ChunkConflictError(BackendAIError, web.HTTPConflict):
+    """
+    Raised when an incoming chunk targets an offset that already holds a
+    different chunk in the upload session (409 Conflict).
+    """
+
+    error_type = "https://api.backend.ai/probs/storage/chunk-conflict"
+    error_title = "Upload Chunk Conflict"
+
+    def error_code(self) -> ErrorCode:
+        return ErrorCode(
+            domain=ErrorDomain.STORAGE_PROXY,
+            operation=ErrorOperation.UPDATE,
+            error_detail=ErrorDetail.CONFLICT,
+        )
+
+
+class UploadSessionCorruptedError(BackendAIError, web.HTTPInternalServerError):
+    """
+    Raised when the on-disk upload session metadata cannot be parsed or is
+    structurally invalid.
+    """
+
+    error_type = "https://api.backend.ai/probs/storage/upload-session-corrupted"
+    error_title = "Upload Session Corrupted"
+
+    def error_code(self) -> ErrorCode:
+        return ErrorCode(
+            domain=ErrorDomain.STORAGE_PROXY,
+            operation=ErrorOperation.READ,
+            error_detail=ErrorDetail.INTERNAL_ERROR,
+        )