AI Studio

A backend platform for running AI training and inference workloads. Submit a training job, the platform schedules it as a containerized task, streams the logs back live, stores the resulting model, and serves it for inference. It is an orchestration platform: the model code is Python in a container, and the Spring Boot backend manages users, queues, scheduling, storage, and APIs around it.

It follows dev/prod parity: the dev profile (the default) runs the same real backends as production — MySQL, Redis, RabbitMQ, MinIO, and Keycloak — from a Docker Compose stack, while the in-JVM stand-ins (mock runner, H2, local storage) are kept for the fast unit-test slice. See docs/DESIGN.md for the architecture, docs/PRODUCTION.md for the hardening plan, and docs/API.md for the HTTP conventions (/api/v1, the response envelope, the error-code catalog, and pagination).

Architecture

A modular monolith: each package under com.aistudio.platform is a module that talks to the others only through service interfaces — the seams for a later microservices split. The core is the job lifecycle, and every piece of infrastructure sits behind an interface with a @ConditionalOnProperty switch, so exactly one implementation is active per profile.

flowchart LR
  U["Client<br/>React console · REST /api/v1"]
  KC["Keycloak<br/>OAuth2 / OIDC"]
  C["Training API<br/>tenant-scoped"]
  DB[("MySQL<br/>job = PENDING + outbox row · one tx")]
  OR["Outbox relay"]
  MQ{{"RabbitMQ"}}
  S["Scheduler<br/>atomic claim · virtual thread · lease"]
  R["ContainerRunner<br/>mock · docker · k8s"]
  SSE["SSE log stream"]
  ST[("MinIO")]
  MR["Model registry<br/>versions · stages"]
  IB["InferenceBackend<br/>in-process · container · k8s"]

  U -->|"POST /training/jobs"| C
  KC -.->|"validate RS256 JWT"| C
  C -->|"one transaction"| DB
  DB -.-> OR -->|publish| MQ --> S
  S --> R
  R -->|stdout| SSE --> U
  R -->|"loss / accuracy"| C
  R -->|artifact| ST
  R -->|on success| MR --> IB
  U -->|"POST /inference/predict"| IB

Loading

Every external dependency is a pluggable seam — the in-JVM variant keeps mvn test dependency-free, while dev/prod default to the real backend:

Seam	Interface	test profile	dev/prod default	cluster
Job execution	ContainerRunner	mock	docker	k8s (one Job per run)
Object storage	StorageService	local	minio	minio
Model serving	InferenceBackend	in-process	container	k8s (Deployment+Service)
Job queue	JobQueue	DB poll	rabbit	rabbit

The scheduler claims work with an atomic UPDATE ... WHERE status = 'PENDING' and carries a heartbeat lease, so job pickup is safe across instances and a crashed worker's job is reclaimed — not double-run. Beyond the core flow the platform adds multi-tenancy (per-org row-level isolation + org-scoped RBAC), a React console (dashboard, training with live logs, datasets, models, experiments, deployments, predict, admin, billing), ML depth (metric ingest + charts, experiments, hyperparameter sweep, lifecycle stages), and a commercial + ops shell (usage metering/billing, notifications, one-command demo, backup/restore, SLO dashboards). See docs/DESIGN.md for the full picture.

Tech stack

• Java 21, Spring Boot 3 (web, security, validation)
• MyBatis with Flyway-managed schema; MySQL 8 (dev/prod), H2 (unit tests)
• Virtual threads for the job scheduler
• docker-java for container orchestration
• MinIO (S3-compatible) for object storage
• OAuth2 Resource Server auth via Keycloak (OIDC), OpenAPI/Swagger UI, SSE for live logs
• Pluggable model serving: one container per model (dev/prod) or in-process (tests)
• Actuator health/readiness + Prometheus metrics (Micrometer), Grafana dashboard
• Pluggable job queue: RabbitMQ (dev/prod default) or DB poll (tests)

Requirements

• JDK 21+
• Maven 3.9+ (mvn). No wrapper is bundled; install Maven, e.g. sudo apt install maven or brew install maven.

Run

The default dev profile targets the real stack, so bring it up first:

docker compose -f deploy/docker-compose.dev.yml up -d   # MySQL, Redis, RabbitMQ, MinIO, Keycloak, observability
mvn spring-boot:run                                     # dev profile → http://localhost:8080

Then open:

• Demo console — http://localhost:8080/
• Swagger UI — http://localhost:8080/swagger-ui.html

Log in with demo / demo123 (or alice / alice123) — the demo accounts live in Keycloak; then submit a job and watch its logs stream. For a fast, dependency-free check, run the unit-test slice instead: mvn test (H2 + in-JVM stand-ins, no Docker).

See it run end to end

One command brings up the whole stack (dependencies + app + bundled console) and waits until it is healthy:

deploy/demo.sh up      # up · down · logs · status

Then walk the full path — submit → train → live curves → register → deploy → predict, plus the commercial layer (metered compute + a job-succeeded notification) — in the console or with the scriptable curl walkthrough in docs/DEMO.md. A real run against the dev stack looks like:

submit demo-run           → job PENDING → RUNNING → SUCCEEDED  (real Docker container)
metric curves             → loss 0.4534 → 0.2467 · accuracy ≈ 0.94  (15 epochs, real SGD)
model auto-registered     → demo-run-model v3
deploy + predict          → prediction 0.9927  (served by the model container)
billing                   → FREE plan · computeSeconds 23 metered
notification              → JOB_SUCCEEDED: demo-run

The console

The bundled React console over that same /api/v1 surface — a training run's live logs and loss/accuracy curves, the model registry with versions + lifecycle stages, and metered billing, all against the real dev stack:

Dashboard	Model registry	Billing

Try it from the command line

# 1. Get an access token from Keycloak (the dev stack's IdP on :8081). The app is an
#    OAuth2 Resource Server: it validates this token via Keycloak's JWKS but issues
#    none of its own. Keycloak also owns registration, refresh, and logout.
TOKEN=$(curl -s -X POST localhost:8081/realms/aistudio/protocol/openid-connect/token \
  -d grant_type=password -d client_id=aistudio-console \
  -d username=demo -d password=demo123 | grep -o '"access_token":"[^"]*' | cut -d'"' -f4)

# (optional) upload a dataset and pass its key as datasetKey when submitting:
# KEY=$(curl -s -X POST localhost:8080/api/v1/datasets -H "Authorization: Bearer $TOKEN" \
#   -F file=@mydata.csv | grep -o '"key":"[^"]*' | cut -d'"' -f4)

# 2. Submit a job
curl -s -X POST localhost:8080/api/v1/training/jobs \
  -H "Authorization: Bearer $TOKEN" -H 'Content-Type: application/json' \
  -d '{"name":"my-job","paramsJson":"{\"epochs\":5}"}'

# 3. Stream its logs (job id 1) — the token goes as a query param (EventSource-friendly)
curl -N "localhost:8080/api/v1/training/jobs/1/logs?token=$TOKEN"

# 4. List models, then run inference
curl -s localhost:8080/api/v1/models -H "Authorization: Bearer $TOKEN"
curl -s -X POST localhost:8080/api/v1/inference/predict \
  -H "Authorization: Bearer $TOKEN" -H 'Content-Type: application/json' \
  -d '{"modelId":1,"features":[0.5,-1.2,3.0]}'

Backend seams & Kubernetes

Real backends are already the dev/prod default (the dev profile targets the Compose stack above; the in-JVM stand-ins are the test default). Each seam can still be set explicitly with an --aistudio.* flag — handy to run one real backend under another profile, or to reach for the Kubernetes implementations below.

Real Docker containers (the dev default) — set explicitly with:

# build the trainer image once
docker build -t aistudio/trainer:latest -f scripts/Dockerfile.trainer scripts
# then run with the docker runner
mvn spring-boot:run -Dspring-boot.run.arguments=--aistudio.runner=docker

Kubernetes Jobs instead of local containers — each training job runs as a K8s Job (fabric8 client; auto-configures from your in-cluster service account or ~/.kube/config). Push the trainer image somewhere the cluster can pull, create the target namespace, and use object storage the cluster can reach (the dataset and artifact travel through it — train.py pulls/pushes by key):

mvn spring-boot:run -Dspring-boot.run.arguments=\
--aistudio.runner=k8s,--aistudio.storage=minio,--aistudio.k8s.namespace=aistudio-jobs

MinIO for storage (the dev default): set aistudio.storage=minio; local writes to the filesystem for tests.

MySQL is the system of record on dev/prod (from the Compose stack); H2 is used only by the test profile. There is no separate mysql profile.

Real model serving instead of the in-process predictor — run each model in its own container:

# build the server image once
docker build -t aistudio/server:latest -f scripts/Dockerfile.server scripts
# then run with the container inference backend
mvn spring-boot:run -Dspring-boot.run.arguments=--aistudio.inference=container

Deploy a trained model and inspect its server (predict also auto-deploys):

curl -s -X POST localhost:8080/api/v1/inference/deployments \
  -H "Authorization: Bearer $TOKEN" -H 'Content-Type: application/json' \
  -d '{"modelId":1}'
curl -s localhost:8080/api/v1/inference/deployments -H "Authorization: Bearer $TOKEN"

RabbitMQ instead of DB polling: start a broker, then run with --aistudio.queue=rabbit. On submit the job id is published to RabbitMQ and a listener claims and runs it; the poll loop stays on as a fallback.

docker run -d --name rabbit -p 5672:5672 rabbitmq:3.13
mvn spring-boot:run -Dspring-boot.run.arguments=--aistudio.queue=rabbit

Run tests

mvn test

Most tests run with the mock runner against H2 — no Docker or MySQL required. The integration tests that need real infrastructure (Docker artifact copy-out / dataset mount, and the MySQL + MinIO Testcontainers tests) detect a Docker daemon and skip automatically when none is reachable, so mvn test stays green anywhere; with Docker present they spin up the real services.

Metrics & health

Spring Boot Actuator is enabled: liveness/readiness probes at /actuator/health (/actuator/health/liveness, /actuator/health/readiness) and a Prometheus scrape endpoint at /actuator/prometheus carrying platform metrics — job counts and duration, queue depth, running jobs, active deployments, inference latency. A ready-to-run Prometheus + Grafana stack lives in deploy/:

cd deploy && docker compose up   # Grafana on :3000, Prometheus on :9090

Project structure

src/main/java/com/aistudio/platform/
  common/       uniform API response + exception handling
  config/       security, docker client, OpenAPI
  security/     tenant context, org-scoped RBAC, rate limit
  user/         Keycloak token -> local user mirror
  org/          organizations, membership, org switcher
  apikey/       per-org, hashed-at-rest programmatic keys
  audit/        append-only audit trail
  storage/      object storage (local / MinIO)
  dataset/      dataset upload + listing
  training/     job lifecycle, REST API, SSE logs, metrics, experiments, sweep
  scheduler/    ContainerRunner (mock/docker/k8s) + the engine, outbox relay
  model/        trained-model registry: versions + lifecycle stages
  inference/    prediction endpoint + InferenceBackend (in-process/container/k8s)
  quota/        per-org quota rollups
  billing/      usage metering + plans (stub/Stripe seam)
  notification/ job-result notifications (log/SMTP seam)
  idempotency/  idempotency-key store
  websocket/    SSE log hub
frontend/       React + TS console (bundled same-origin behind -Pfrontend)
scripts/        example Python trainer + model server + Dockerfiles
deploy/         Compose dev stack, Helm chart, observability, backups, demo.sh
docs/           DESIGN · PRODUCTION · ROADMAP · API · DEMO · OPERATIONS · SECURITY

Roadmap

See docs/ROADMAP.md for the prioritized plan and docs/DESIGN.md for the architecture and the detailed Phase 4 design.

License

MIT © 2026 jianglulu