AmineDiro

use flume;
use futures::future::join_all;
use futures::stream::iter;
use futures_util::stream::FuturesUnordered;
use futures_util::{SinkExt, StreamExt};
use std::borrow::Cow;
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use test_websockets::ClientMessage;
use tokio::fs::OpenOptions;
use tokio::io::{AsyncWriteExt, BufWriter};

AmineDiro

use axum::{
    extract::ws::{Message, WebSocket, WebSocketUpgrade},
    response::IntoResponse,
    routing::get,
    Router,
};
use test_websockets::ClientMessage;
use tracing::{debug, error, info};

use std::net::SocketAddr;

AmineDiro

[package]
name = "uploader"
version = "0.1.0"
edition = "2021"

[dependencies]
anyhow = "1.0.86"
clap = { version = "4.5.13", features = ["derive"] }
futures = "0.3.30"
reqwest = { version = "0.12.5", features = ["multipart", "stream"] }

AmineDiro

CREATE TYPE job_status AS ENUM ('pending', 'in_progress', 'done', 'failed');

CREATE TABLE jobs (
    id SERIAL PRIMARY KEY,
    status job_status NOT NULL DEFAULT 'pending',
    payload JSONB,
    created_at TIMESTAMP DEFAULT now(),
    updated_at TIMESTAMP DEFAULT now(),
    visible_at TIMESTAMP DEFAULT now(), -- SQS visibility timeout
    retry_count INT DEFAULT 0

AmineDiro

# This isn't supposed to run as a bash script, i named it with ".sh" for syntax highlighting.

# https://developer.nvidia.com/nsight-systems
# https://docs.nvidia.com/nsight-systems/profiling/index.html

# My preferred nsys (command line executable used to create profiles) commands
#
# In your script, write
# torch.cuda.nvtx.range_push("region name")
# ...

AmineDiro

"""
The most atomic way to train and inference a GPT in pure, dependency-free Python.
This file is the complete algorithm.
Everything else is just efficiency.

@karpathy
"""

import os       # os.path.exists
import math     # math.log, math.exp

AmineDiro

#!/usr/bin/env -S uv run
# /// script
# requires-python = ">=3.10, <3.13"
# dependencies = [
#     "aiohttp==3.13.3",
#     "datasets==4.6.1",
#     "hf-transfer==0.1.9",
#     "requests==2.32.5",
#     "torch==2.10.0",
#     "transformers==4.57.6",

AmineDiro

#!/usr/bin/env -S uv run
# /// script
# requires-python = ">=3.10, <3.13"
# dependencies = [
#     "aiohttp==3.13.3",
#     "datasets==4.6.1",
#     "hf-transfer==0.1.9",
#     "requests==2.32.5",
#     "torch==2.10.0",
#     "transformers==4.57.6",

AmineDiro

"""
Minimal sanity-check for AsyncGRPOTrainer: the "Immediate EOS" test.

The model is rewarded with R(y) = -len(completion_tokens). The optimal policy
is to emit <EOS> immediately (reward = -1). Within a handful of steps the
average completion length should drop and reward_mean should climb toward -1.

Start the vLLM server:

    CUDA_VISIBLE_DEVICES=0 VLLM_SERVER_DEV_MODE=1 vllm serve Qwen/Qwen3-0.6B \

AmineDiro

"""

HOW TO RUN
-----------
Step 1 – Start a vLLM server with data-parallel support (replace N with the number of DP shards,
          and adjust --tensor-parallel-size / --gpu-memory-utilization as needed):

CUDA_VISIBLE_DEVICES=2,3,4,5 VLLM_SERVER_DEV_MODE=1 vllm serve Qwen/Qwen3-4B \
    --data-parallel-size 4 \
    --tensor-parallel-size 1 \