diff --git a/backend/src/backend/main.py b/backend/src/backend/main.py index b8b474c..61fcc40 100644 --- a/backend/src/backend/main.py +++ b/backend/src/backend/main.py @@ -34,11 +34,18 @@ async def lifespan(app: Any) -> AsyncIterator[None]: access_key=settings.rustfs_access_key, secret_key=settings.rustfs_secret_key, ) + # Ensure all three purpose-named buckets exist; the storage edge picks + # the right one per upload (see resolve_bucket in storage_api.py). + for bucket in ( + settings.rustfs_workspace_bucket, + settings.rustfs_version_bucket, + settings.rustfs_run_log_bucket, + ): + await asyncio.to_thread( + app.state.object_store.ensure_bucket, + bucket, + ) app.state.default_bucket = settings.rustfs_workspace_bucket - await asyncio.to_thread( - app.state.object_store.ensure_bucket, - app.state.default_bucket, - ) storage_http_client = httpx.AsyncClient( transport=httpx.ASGITransport(app=app), base_url="http://backend.internal", diff --git a/backend/src/backend/storage_api.py b/backend/src/backend/storage_api.py index c49261b..2447b79 100644 --- a/backend/src/backend/storage_api.py +++ b/backend/src/backend/storage_api.py @@ -58,6 +58,39 @@ def safe_file_name(value: str) -> str: return name +# Map an upload's usage_type to the RustFS bucket that should hold the +# resulting object. ``usage_type`` is the only signal available at the +# storage edge (the request comes from either the public API or the +# internal schedule worker), so we make the routing decision in one place +# here and let every other layer — server-object create, multipart upload, +# direct put — inherit the mapping. +BUCKET_FOR_USAGE: dict[str, str] = { + "working_copy": settings.rustfs_workspace_bucket, + "public_script": settings.rustfs_workspace_bucket, + "data_resource": settings.rustfs_workspace_bucket, + "snapshot": settings.rustfs_workspace_bucket, + "version_artifact": settings.rustfs_version_bucket, + "run_log": settings.rustfs_run_log_bucket, + "run_result": settings.rustfs_run_log_bucket, +} + + +def resolve_bucket( + usage_type: str, + *, + workspace: Workspaces, +) -> str: + """Pick the bucket for ``usage_type``. + + ``workspace.artifact_bucket`` (per-workspace override) wins over the + usage-type default. An unknown ``usage_type`` falls through to the + workspace bucket so we never silently drop an object. + """ + if workspace.artifact_bucket: + return workspace.artifact_bucket + return BUCKET_FOR_USAGE.get(usage_type, settings.rustfs_workspace_bucket) + + def storage_payload(item: StorageObjects) -> dict[str, Any]: return { "storage_object_id": item.storage_object_id, @@ -165,9 +198,7 @@ async def create_upload_record( upload = existing else: upload_id = new_ulid() - bucket_name = ( - workspace.artifact_bucket or request.app.state.default_bucket - ) + bucket_name = resolve_bucket(payload.usage_type, workspace=workspace) # Object key is a flat two-level path: workspace id + upload id. The # original file name and content type live in the StorageObjects row # (file_name / mime_type / object_key) — they are not part of the diff --git a/common/src/common/config.py b/common/src/common/config.py new file mode 100644 index 0000000..f00c121 --- /dev/null +++ b/common/src/common/config.py @@ -0,0 +1,141 @@ +"""Centralised configuration via pydantic-settings. + +All Python services (backend / schedule / runtime) import ``settings`` from +this module. The instance is a process-wide singleton (lru_cache wrapped), +so each field is parsed from the environment exactly once at first access. + +Rules for adding a new variable: + + 1. Add the field here with a sensible default that lets local dev + boot without the env var set. + 2. Use ``settings.`` at the call site. Never re-introduce + ``os.environ`` / ``os.getenv`` for application config — they + bypass this central registry. + 3. Document the env var name in ``.env.example`` so operators know + it exists. +""" + +from __future__ import annotations + +from functools import lru_cache + +from pydantic import Field +from pydantic_settings import BaseSettings, SettingsConfigDict + + +class Settings(BaseSettings): + # ── database ────────────────────────────────────────────────── + database_url: str = Field( + default=( + "mysql+asyncmy://model_platform:model_platform@mysql:3306/" + "model_platform?charset=utf8mb4" + ), + description="SQLAlchemy async URI for the platform MySQL.", + ) + + # ── JWT ─────────────────────────────────────────────────────── + jwt_secret: str = Field( + default="dev-only-not-for-production", + description="HS256 secret used by backend's jupyter auth_request.", + ) + + # ── service identity ────────────────────────────────────────── + service_name: str = Field( + default="service", + description="Service label surfaced in lifespan / health checks.", + ) + + # ── runtime container endpoint ─────────────────────────────── + runtime_api_url: str = Field( + default="http://runtime:8000", + description="Backend → Runtime HTTP endpoint.", + ) + + # ── RustFS object storage ──────────────────────────────────── + rustfs_endpoint: str = Field( + default="http://rustfs:9000", + description="S3 endpoint for the RustFS upstream.", + ) + rustfs_access_key: str = Field( + default="modelplatform", + description="boto3 access key for RustFS.", + ) + rustfs_secret_key: str = Field( + default="modelplatformsecret", + description="boto3 secret key for RustFS.", + ) + rustfs_workspace_bucket: str = Field( + default="workspaces", + description=( + "Bucket for workspace files (notebooks / scripts / working " + "copies). Layout: s3:////..." + ), + ) + rustfs_version_bucket: str = Field( + default="versions", + description="Bucket for immutable script-version artifacts.", + ) + rustfs_run_log_bucket: str = Field( + default="run-logs", + description="Bucket for schedule run logs.", + ) + + # ── local FS roots ──────────────────────────────────────────── + workspace_root: str = Field( + default="/workspace/workspaces", + description=( + "Schedule subprocess cwd; staging area for notebook_runner. " + "Backend no longer writes here — workspace files live in " + "RustFS via rustfs_workspace_bucket." + ), + ) + workspaces_root: str = Field( + default="/app/workspaces", + description="Runtime rclone FUSE mount point for the workspace bucket.", + ) + + # ── rclone remote spec ──────────────────────────────────────── + remote_bucket: str = Field( + default="rustfs:workspaces", + description="rclone remote spec for the workspace bucket.", + ) + + # ── schedule → backend API ──────────────────────────────────── + backend_api_url: str = Field( + default="http://backend:8000", + description="Schedule → Backend HTTP base URL (cron post-back).", + ) + + # ── runtime public base URL ────────────────────────────────── + public_base_url: str = Field( + default="http://runtime", + description="Public base URL for the runtime container.", + ) + + # ── readiness probes ────────────────────────────────────────── + readiness_targets: str = Field( + default="", + description=( + "Comma-separated host:port list checked by /health/ready. " + "Empty disables the check." + ), + ) + + model_config = SettingsConfigDict( + env_file=".env", + env_file_encoding="utf-8", + case_sensitive=False, + extra="ignore", + ) + + +@lru_cache +def get_settings() -> Settings: + """Construct (and cache) the singleton Settings instance.""" + return Settings() + + +settings: Settings = get_settings() + + +__all__ = ["Settings", "get_settings", "settings"] \ No newline at end of file diff --git a/common/src/common/storage/client.py b/common/src/common/storage/client.py new file mode 100644 index 0000000..dc6b496 --- /dev/null +++ b/common/src/common/storage/client.py @@ -0,0 +1,144 @@ +"""HTTP client for the internal ``/internal/v1/...`` storage surface. + +The class accepts an injected ``httpx.AsyncClient`` so callers can wire +it up with either a real network transport or an ``ASGITransport`` for +in-process dispatch. ``SchedulerStorageClient`` extends this class to +add a single ``create_object`` helper used by the schedule service. + +All public methods raise :class:`StorageClientError` (or a subclass) on +failure. Translation to a web-framework exception (e.g. FastAPI's +``HTTPException``) is the caller's responsibility so this client stays +usable from non-FastAPI contexts such as the schedule worker. +""" + +from __future__ import annotations + +import base64 +from typing import Any + +import httpx + + +class StorageClientError(Exception): + """Base class for storage client failures.""" + + +class StorageUnavailable(StorageClientError): + """Transport-level failure; safe to retry.""" + + def __init__(self, message: str) -> None: + super().__init__(message) + self.retryable = True + + +class StorageRequestFailed(StorageClientError): + """Storage endpoint returned a non-2xx response.""" + + def __init__(self, status_code: int, detail: Any) -> None: + super().__init__(f"storage request failed with status {status_code}") + self.status_code = status_code + self.detail = detail + + +class StorageClient: + def __init__(self, client: httpx.AsyncClient) -> None: + self.client = client + + async def _request( + self, + method: str, + path: str, + *, + payload: dict[str, Any] | None = None, + ) -> dict[str, Any]: + try: + response = await self.client.request( + method, + path, + json=payload, + ) + except httpx.RequestError as exc: + raise StorageUnavailable("Storage service temporarily unavailable") from exc + if response.is_error: + try: + detail = response.json().get("detail", response.text) + except ValueError: + detail = response.text + raise StorageRequestFailed(response.status_code, detail) + return response.json() + + async def create_upload( + self, + payload: dict[str, Any], + ) -> dict[str, Any]: + return (await self._request( + "POST", + "/internal/v1/uploads", + payload=payload, + ))["data"] + + async def complete_upload( + self, + upload_id: str, + payload: dict[str, Any], + ) -> dict[str, Any]: + return (await self._request( + "POST", + f"/internal/v1/uploads/{upload_id}/complete", + payload=payload, + ))["data"] + + async def create_server_object( + self, + *, + workspace_id: str, + user_id: str, + usage_type: str, + file_name: str, + content_type: str, + content: bytes, + visibility: str, + is_immutable: bool, + idempotency_key: str, + ) -> dict[str, Any]: + result = await self._request( + "POST", + "/internal/v1/objects", + payload={ + "workspace_id": workspace_id, + "user_id": user_id, + "usage_type": usage_type, + "file_name": file_name, + "content_type": content_type, + "content_base64": base64.b64encode(content).decode("ascii"), + "visibility": visibility, + "is_immutable": is_immutable, + "idempotency_key": idempotency_key, + }, + ) + return result["data"] + + async def create_download_url( + self, + storage_object_id: str, + expires_seconds: int, + ) -> dict[str, Any]: + return (await self._request( + "POST", + f"/internal/v1/objects/{storage_object_id}/download-url", + payload={"expires_seconds": expires_seconds}, + ))["data"] + + async def delete_object(self, storage_object_id: str) -> dict[str, Any]: + return (await self._request( + "DELETE", + f"/internal/v1/objects/{storage_object_id}", + ))["data"] + + +__all__ = [ + "StorageClient", + "StorageClientError", + "StorageRequestFailed", + "StorageUnavailable", +] diff --git a/common/src/common/storage/schemas.py b/common/src/common/storage/schemas.py new file mode 100644 index 0000000..e491774 --- /dev/null +++ b/common/src/common/storage/schemas.py @@ -0,0 +1,84 @@ +"""Pydantic request models for the internal ``/internal/v1/...`` storage surface.""" + +from __future__ import annotations + +from typing import Literal + +from pydantic import Field, field_validator + +from common.schemas import StrictModel + + +__all__ = [ + "CompleteUploadRequest", + "CreateUploadRequest", + "DownloadUrlRequest", + "ServerObjectRequest", +] + + +class CreateUploadRequest(StrictModel): + workspace_id: str = Field(min_length=26, max_length=26) + user_id: str = Field(min_length=26, max_length=26) + usage_type: Literal[ + "data_resource", + "version_artifact", + "snapshot", + "run_log", + "run_result", + "working_copy", + "public_script", + ] + file_name: str = Field(min_length=1, max_length=255) + content_type: str = Field(min_length=1, max_length=255) + expected_size_bytes: int = Field(ge=0, le=100 * 1024 * 1024) + expected_hash: str | None = Field(default=None, min_length=64, max_length=64) + idempotency_key: str = Field(min_length=8, max_length=128) + + @field_validator("expected_hash") + @classmethod + def validate_hash(cls, value: str | None) -> str | None: + if value is None: + return None + normalized = value.lower() + if any(character not in "0123456789abcdef" for character in normalized): + raise ValueError("expected_hash must be lowercase SHA-256 hex") + return normalized + + +class CompleteUploadRequest(StrictModel): + usage_type: Literal[ + "data_resource", + "version_artifact", + "snapshot", + "run_log", + "run_result", + "working_copy", + "public_script", + ] + visibility: Literal["private", "workspace", "public"] = "private" + is_immutable: bool = False + + +class ServerObjectRequest(StrictModel): + workspace_id: str = Field(min_length=26, max_length=26) + user_id: str = Field(min_length=26, max_length=26) + usage_type: Literal[ + "data_resource", + "version_artifact", + "snapshot", + "run_log", + "run_result", + "working_copy", + "public_script", + ] + file_name: str = Field(min_length=1, max_length=255) + content_type: str = Field(min_length=1, max_length=255) + content_base64: str = Field(min_length=1) + visibility: Literal["private", "workspace", "public"] = "private" + is_immutable: bool = False + idempotency_key: str = Field(min_length=8, max_length=128) + + +class DownloadUrlRequest(StrictModel): + expires_seconds: int = Field(default=300, ge=30, le=3600)