mbpp-grpo-pipeline-test-10-samples.py

{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"},"kaggle":{"accelerator":"tpuV5e8","dataSources":[],"dockerImageVersionId":31236,"isInternetEnabled":true,"language":"python","sourceType":"script","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# %% [code]\n\"\"\"\nMinimal test to verify Kaggle environment setup with uv.\n\"\"\"\nimport sys\nimport subprocess\nimport os\n\nprint(\"=\" * 80)\nprint(\"Kaggle Environment Test with uv\")\nprint(\"=\" * 80)\n\n# Step 0: Install uv and sync dependencies\nprint(\"\\n[0/5] Installing uv and syncing dependencies...\")\ntry:\n    # Install uv\n    print(\"  Installing uv...\")\n    result = subprocess.run(\n        [\"pip\", \"install\", \"-q\", \"uv\"],\n        capture_output=True,\n        text=True,\n        timeout=60\n    )\n    if result.returncode == 0:\n        print(\"  ✓ uv installed\")\n    else:\n        print(f\"  ✗ Failed to install uv: {result.stderr}\")\n        sys.exit(1)\n\n    # Sync dependencies using uv\n    print(\"  Running uv sync...\")\n    result = subprocess.run(\n        [\"uv\", \"sync\"],\n        capture_output=True,\n        text=True,\n        timeout=300\n    )\n    if result.returncode == 0:\n        print(\"  ✓ Dependencies synced successfully\")\n        print(f\"  Output: {result.stdout[:200]}...\")\n    else:\n        print(f\"  ✗ uv sync failed: {result.stderr[:500]}\")\n        # Continue anyway to see what's available\n\nexcept Exception as e:\n    print(f\"  ✗ Error during uv setup: {e}\")\n    # Continue anyway\n\n# Step 1: Check Python version\nprint(f\"\\n[1/5] Python version: {sys.version}\")\n\n# Step 2: Check GPU availability\nprint(\"\\n[2/5] Checking GPU...\")\ntry:\n    import jax\n    print(f\"  JAX version: {jax.__version__}\")\n    devices = jax.devices()\n    print(f\"  Devices: {devices}\")\n    print(f\"  Device type: {devices[0].platform}\")\n    print(f\"  Number of devices: {len(devices)}\")\nexcept Exception as e:\n    print(f\"  Error: {e}\")\n\n# Step 3: Check if we can import key libraries\nprint(\"\\n[3/5] Checking imports...\")\ntry:\n    import optax\n    print(f\"  ✓ optax {optax.__version__}\")\nexcept Exception as e:\n    print(f\"  ✗ optax: {e}\")\n\ntry:\n    from flax import nnx\n    print(f\"  ✓ flax.nnx\")\nexcept Exception as e:\n    print(f\"  ✗ flax.nnx: {e}\")\n\ntry:\n    from huggingface_hub import snapshot_download\n    print(f\"  ✓ huggingface_hub\")\nexcept Exception as e:\n    print(f\"  ✗ huggingface_hub: {e}\")\n\n# Step 4: Check if tunix is available\nprint(\"\\n[4/5] Checking tunix library...\")\ntry:\n    from tunix.rl import rl_cluster as rl_cluster_lib\n    print(f\"  ✓ tunix.rl\")\nexcept Exception as e:\n    print(f\"  ✗ tunix.rl: {e}\")\n    print(f\"  Installing tunix...\")\n    import subprocess\n    result = subprocess.run(\n        [\"pip\", \"install\", \"-q\", \"git+https://github.com/google-deepmind/tunix.git\"],\n        capture_output=True,\n        text=True\n    )\n    if result.returncode == 0:\n        print(f\"  ✓ tunix installed successfully\")\n    else:\n        print(f\"  ✗ Failed to install tunix: {result.stderr}\")\n\n# Step 5: Test complete\nprint(\"\\n[5/5] Environment check complete!\")\nprint(\"=\" * 80)\nprint(\"✅ Ready to run GRPO training pipeline\")\nprint(\"=\" * 80)\n","metadata":{"_uuid":"efa5a11f-0a20-4ecc-b8c8-1bcbe061bf5a","_cell_guid":"e4739659-1d06-42ec-87b8-96319961bf18","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-12-25T20:47:26.451164Z","iopub.execute_input":"2025-12-25T20:47:26.451432Z","iopub.status.idle":"2025-12-25T20:47:27.394031Z","shell.execute_reply.started":"2025-12-25T20:47:26.451413Z","shell.execute_reply":"2025-12-25T20:47:27.392702Z"}},"outputs":[{"name":"stdout","text":"================================================================================\nKaggle Environment Test with uv\n================================================================================\n\n[0/5] Installing uv and syncing dependencies...\n  Installing uv...\n  ✓ uv installed\n  Running uv sync...\n  ✓ Dependencies synced successfully\n  Output: ...\n\n[1/5] Python version: 3.12.12 (main, Dec  9 2025, 02:04:51) [GCC 14.2.0]\n\n[2/5] Checking GPU...\n  JAX version: 0.8.1\n  Devices: [TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0), TpuDevice(id=1, process_index=0, coords=(1,0,0), core_on_chip=0), TpuDevice(id=2, process_index=0, coords=(0,1,0), core_on_chip=0), TpuDevice(id=3, process_index=0, coords=(1,1,0), core_on_chip=0), TpuDevice(id=4, process_index=0, coords=(0,2,0), core_on_chip=0), TpuDevice(id=5, process_index=0, coords=(1,2,0), core_on_chip=0), TpuDevice(id=6, process_index=0, coords=(0,3,0), core_on_chip=0), TpuDevice(id=7, process_index=0, coords=(1,3,0), core_on_chip=0)]\n  Device type: tpu\n  Number of devices: 8\n\n[3/5] Checking imports...\n  ✓ optax 0.2.6\n  ✓ flax.nnx\n  ✓ huggingface_hub\n\n[4/5] Checking tunix library...\n  ✓ tunix.rl\n\n[5/5] Environment check complete!\n================================================================================\n✅ Ready to run GRPO training pipeline\n================================================================================\n","output_type":"stream"}],"execution_count":21}]}