{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Utility notebook to get the storage per project and user\n",
    "\n",
    "<a target=\"_blank\" href=\"https://colab.research.google.com/github/neptune-ai/examples/blob/main/utils/management_tools/Get_storage_used_by_project_user.ipynb\">\n",
    "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open in Colab\"/>\n",
    "</a><a target=\"_blank\" href=\"https://github.com/neptune-ai/examples/blob/main/utils/management_tools/Get_storage_used_by_project_user.ipynb\">\n",
    "  <img alt=\"Open in GitHub\" src=\"https://img.shields.io/badge/Open_in_GitHub-blue?logo=github&labelColor=black\">\n",
    "</a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install -qq -U neptune tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import neptune\n",
    "import pandas as pd\n",
    "from neptune import management\n",
    "from tqdm.auto import tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set Neptune API token"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if \"NEPTUNE_API_TOKEN\" not in os.environ:\n",
    "    from getpass import getpass\n",
    "\n",
    "    os.environ[\"NEPTUNE_API_TOKEN\"] = getpass(\"Enter the Neptune API token you wish to use: \")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Enter workspace to scan"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "workspace = input(\"Enter workspace name:\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get workspace status\n",
    "\n",
    "Shows the storage limit and storage used, number of active projects, active project limit, and number of members in the workspace."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "management.get_workspace_status(workspace)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## List all projects in the workspace\n",
    "This will list all the projects in the workspace accessible by the NEPTUNE_API_TOKEN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "projects = [\n",
    "    project for project in management.get_project_list() if project.split(\"/\")[0] == workspace\n",
    "]\n",
    "projects"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get the runs, models, and model versions tables for all the projects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_runs_df = all_models_df = all_model_versions_df = pd.DataFrame()\n",
    "\n",
    "for project in tqdm(projects):\n",
    "    with neptune.init_project(project=project, mode=\"read-only\") as proj:\n",
    "        runs_df = proj.fetch_runs_table(\n",
    "            columns=[\"sys/owner\", \"sys/size\"],\n",
    "            progress_bar=None,\n",
    "        ).to_pandas()\n",
    "\n",
    "        models_df = proj.fetch_models_table(\n",
    "            columns=[\"sys/owner\", \"sys/size\"],\n",
    "            progress_bar=None,\n",
    "        ).to_pandas()\n",
    "\n",
    "    if not runs_df.empty:\n",
    "        runs_df = runs_df.sort_values(by=[\"sys/size\"], ascending=False)\n",
    "        runs_df[\"project\"] = project\n",
    "\n",
    "        all_runs_df = pd.concat([all_runs_df, runs_df], ignore_index=True)\n",
    "\n",
    "    if not models_df.empty:\n",
    "        models_df = models_df.sort_values(by=[\"sys/size\"], ascending=False)\n",
    "        models_df[\"project\"] = project\n",
    "\n",
    "        all_models_df = pd.concat([all_models_df, models_df], ignore_index=True)\n",
    "\n",
    "        for model_id in tqdm(models_df[\"sys/id\"]):\n",
    "            with neptune.init_model(project=project, with_id=model_id, mode=\"read-only\") as model:\n",
    "                model_versions_df = model.fetch_model_versions_table(\n",
    "                    columns=[\"sys/owner\", \"sys/size\"],\n",
    "                    progress_bar=None,\n",
    "                ).to_pandas()\n",
    "\n",
    "                if not model_versions_df.empty:\n",
    "                    model_versions_df = model_versions_df.sort_values(\n",
    "                        by=[\"sys/size\"], ascending=False\n",
    "                    )\n",
    "\n",
    "                    all_model_versions_df = pd.concat(\n",
    "                        [all_model_versions_df, model_versions_df], ignore_index=True\n",
    "                    )\n",
    "\n",
    "                all_model_versions_df[\"project\"] = project"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_runs_df[\"type\"] = \"run\"\n",
    "all_models_df[\"type\"] = \"model\"\n",
    "all_model_versions_df[\"type\"] = \"model_version\"\n",
    "\n",
    "all_objects_df = pd.concat([all_runs_df, all_models_df, all_model_versions_df], ignore_index=True)\n",
    "all_objects_df.sort_values(by=[\"sys/size\"], ascending=False, inplace=True)\n",
    "all_objects_df.reset_index(drop=True, inplace=True)\n",
    "\n",
    "all_objects_df[\"size_gb\"] = all_objects_df[\"sys/size\"] / (1024**3)\n",
    "all_objects_df.drop(\"sys/size\", axis=1, inplace=True)\n",
    "all_objects_df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Storage per project"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_objects_df.groupby(\"project\", as_index=False)[[\"size_gb\"]].sum().sort_values(\n",
    "    \"size_gb\", ascending=False\n",
    ").reset_index(drop=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Storage per user"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_objects_df.groupby(\"sys/owner\", as_index=False)[[\"size_gb\"]].sum().sort_values(\n",
    "    \"size_gb\", ascending=False\n",
    ").reset_index(drop=True)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "private_outputs": true,
   "provenance": [],
   "toc_visible": true
  },
  "gpuClass": "standard",
  "kernelspec": {
   "display_name": "py310",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}