{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "360bb8aa",
   "metadata": {},
   "source": [
    "## Run as script"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "203322a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# python /home/georg/notebooks/web_harvest/harvest/podcasts/02_quantize_data_script.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "659dab7d",
   "metadata": {},
   "source": [
    "## run in notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "e9f10789",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import gc\n",
    "import json\n",
    "import random\n",
    "import numpy as np\n",
    "import tempfile\n",
    "import time\n",
    "import funcy\n",
    "import tqdm\n",
    "import shutil\n",
    "import multiprocessing\n",
    "\n",
    "from suno_utils.utils.text import read_jsonl\n",
    "from suno_utils.web.harvest import get_file_ext, get_filename\n",
    "\n",
    "DATA_DIR = \"/home/georg/notebooks/web_harvest/harvest/podcasts/data\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "52bf4241",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "414,537 podcasts harvested\n"
     ]
    }
   ],
   "source": [
    "metas = read_jsonl(os.path.join(DATA_DIR, \"01_harvest.jsonl\"))\n",
    "podcast_uid_set = set([m[\"id\"] for m in metas if m[\"success\"]])\n",
    "print(f\"{len(podcast_uid_set):,}\", \"podcasts harvested\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "b78326b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # load poscast meta\n",
    "# from suno_utils.utils.podcasts import load_podcast_db\n",
    "# meta_df = load_podcast_db(\"data/podcastindex_feeds.db\", english_only=True, anchor_only=True)\n",
    "# meta_df = meta_df[meta_df[\"uid\"].isin(podcast_uid_set)].reset_index(drop=True)\n",
    "# assert(meta_df.shape[0] == len(podcast_uid_set))\n",
    "# print(f\"{meta_df[meta_df['category1'] == 'music'].shape[0]:,}\", \"music podcasts\")\n",
    "# # 11,680 music podcasts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "f3615ef3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "414,419 podcasts collected\n",
      "2,061,387 episodes collected\n",
      "1,030,694 hours of data\n"
     ]
    }
   ],
   "source": [
    "# collect episodes for each working uuid\n",
    "# TODO: exclude things like music?\n",
    "n_podcasts = 0\n",
    "n_episodes = 0\n",
    "work_items = []\n",
    "for m in metas:\n",
    "    if not m[\"success\"]:\n",
    "        continue\n",
    "    tmp_n_episodes = 0\n",
    "    for filename in m[\"meta\"].get(\"audio_filenames\", []):\n",
    "        podcast_uid = m[\"id\"]\n",
    "        episode_uid = get_filename(filename)\n",
    "        rel_filepath = os.path.join(podcast_uid, filename)\n",
    "        work_items.append((podcast_uid, episode_uid, rel_filepath))\n",
    "        tmp_n_episodes += 1\n",
    "        n_episodes += 1\n",
    "    if tmp_n_episodes > 0:\n",
    "        n_podcasts += 1\n",
    "random.seed(6006)\n",
    "random.shuffle(work_items)\n",
    "print(f\"{n_podcasts:,}\", \"podcasts collected\")\n",
    "print(f\"{n_episodes:,}\", \"episodes collected\")\n",
    "print(f\"{int(round(n_episodes * 0.5)):,}\", \"hours of data\")\n",
    "# 414,419 podcasts collected\n",
    "# 2,061,387 episodes collected\n",
    "# 1,030,694 hours of data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad219452",
   "metadata": {},
   "source": [
    "## Process data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "81b9f479",
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "\n",
    "from suno_utils.tasks.encodec import encode, load_audio_mp\n",
    "\n",
    "S3_BUCKET = \"suno-data\"\n",
    "S3_DIR = \"datasets/harvest/podcasts\"\n",
    "\n",
    "S3_RAW_AUDIO_DIR = os.path.join(S3_DIR, \"audio\")\n",
    "TO_S3_ARRAY_DIR = os.path.join(S3_DIR, \"hubert_embeddings\")\n",
    "\n",
    "MIN_DURATION_S = 2*60\n",
    "MAX_DURATION_S = 1*60*60\n",
    "\n",
    "SAMPLE_RATE = 24_000\n",
    "\n",
    "MIN_ARRAY_LEN = int(MIN_DURATION_S * SAMPLE_RATE)\n",
    "MAX_ARRAY_LEN = int(MAX_DURATION_S * SAMPLE_RATE)\n",
    "\n",
    "def s3_download_audio(s3_work_item):\n",
    "    try:\n",
    "        s3_bucket = boto3.resource(\"s3\").Bucket(S3_BUCKET)\n",
    "        tmp_dir, uid, filepath = s3_work_item\n",
    "        from_fp = os.path.join(S3_RAW_AUDIO_DIR, filepath)\n",
    "        to_fp = os.path.join(tmp_dir, f\"{uid}.{get_file_ext(filepath)}\")\n",
    "        s3_bucket.download_file(from_fp, to_fp)\n",
    "    except:\n",
    "        return None\n",
    "    return uid, to_fp\n",
    "\n",
    "def s3_load_audios(uids, filepaths, n_cores=10, num_workers_loader=32, prefetch_factor_loader=2):\n",
    "    assert(len(uids) == len(filepaths))\n",
    "    with tempfile.TemporaryDirectory() as tmp_dir:\n",
    "        s3_work_items = [\n",
    "            (tmp_dir, uid, filepath) \n",
    "            for tmp_dir, uid, filepath in zip(\n",
    "                [tmp_dir] * len(uids), uids, filepaths\n",
    "            )\n",
    "        ]\n",
    "        with multiprocessing.Pool(n_cores) as p:\n",
    "            out = p.map(s3_download_audio, s3_work_items)\n",
    "            out = [e for e in out if e is not None]\n",
    "        local_uids, local_filepaths = zip(*out)\n",
    "        audio_array_list = load_audio_mp(\n",
    "            local_filepaths, \n",
    "            max_array_len=MAX_ARRAY_LEN, \n",
    "            num_workers=num_workers_loader, \n",
    "            prefetch_factor=prefetch_factor_loader,\n",
    "        )  # returns [(1, n_samples)]\n",
    "    # remove unsuitable arrays\n",
    "    indexed_audio_array_list = [\n",
    "        (uid, arr) \n",
    "        for uid, arr in zip(local_uids, audio_array_list) \n",
    "        if arr is not None and arr.shape[-1] >= MIN_ARRAY_LEN\n",
    "    ]\n",
    "    return indexed_audio_array_list\n",
    "\n",
    "def create_data(\n",
    "    work_items, \n",
    "    out_meta_filepath, \n",
    "    allow_continue=True, \n",
    "    n_gpus=4,\n",
    "    chunksize=100, \n",
    "    gpu_batch_size=16,\n",
    "    n_cores=10,\n",
    "    num_workers_loader=32,\n",
    "    prefetch_factor_loader=2,\n",
    "):\n",
    "    array_metas = []\n",
    "    if allow_continue and os.path.exists(out_meta_filepath):\n",
    "        s3_bucket = boto3.resource(\"s3\").Bucket(S3_BUCKET)\n",
    "        n_processed = 0\n",
    "        for e in s3_bucket.objects.filter(Prefix=TO_S3_ARRAY_DIR):\n",
    "            if e.key.endswith(\".npz\"):\n",
    "                n_processed += 1\n",
    "        with open(out_meta_filepath) as f:\n",
    "            array_metas = json.load(f)\n",
    "        assert(len(array_metas) == n_processed)\n",
    "        print(n_processed, \"chunks already processed\")\n",
    "    # run loop\n",
    "    n_chunk = len(array_metas)\n",
    "    work_items_offset = int(n_chunk * chunksize)\n",
    "    n_steps_remaining = int(np.ceil(len(work_items[work_items_offset:]) / chunksize))\n",
    "    for work_items_chunk in tqdm.tqdm(\n",
    "        funcy.chunks(chunksize, work_items[work_items_offset:]), total=n_steps_remaining\n",
    "    ):\n",
    "        podcast_uids, episode_uids, s3_filepaths = zip(*work_items_chunk)\n",
    "        # load audios\n",
    "        indexed_audio_array_list = s3_load_audios(\n",
    "            episode_uids, s3_filepaths, \n",
    "            n_cores=n_cores, \n",
    "            num_workers_loader=num_workers_loader, \n",
    "            prefetch_factor_loader=prefetch_factor_loader,\n",
    "        )\n",
    "        # ecode audios\n",
    "        retained_uids, audio_array_list = zip(*indexed_audio_array_list)\n",
    "        encoded_array_list = encode(audio_array_list, batch_size=gpu_batch_size, bandwidth=12, n_gpus=n_gpus)\n",
    "        del audio_array_list, indexed_audio_array_list\n",
    "        # save as numpy archive\n",
    "        encoded_array_map = {uid: arr for uid, arr in zip(retained_uids, encoded_array_list)}\n",
    "        del encoded_array_list\n",
    "        # save on s3\n",
    "        out_array_filename = f\"part_{n_chunk}.npz\"\n",
    "        s3_bucket = boto3.resource(\"s3\").Bucket(S3_BUCKET)\n",
    "        out_s3_array_filepath = os.path.join(TO_S3_ARRAY_DIR, out_array_filename)\n",
    "        with tempfile.TemporaryDirectory() as tmp_dir:\n",
    "            out_encoded_arrays_filepath = os.path.join(tmp_dir, out_array_filename)\n",
    "            np.savez(out_encoded_arrays_filepath[:-4], **encoded_array_map)\n",
    "            s3_bucket.upload_file(out_encoded_arrays_filepath, out_s3_array_filepath)\n",
    "        # write meta data file\n",
    "        array_metas.append({\n",
    "            \"s3_filepath\": out_s3_array_filepath,\n",
    "            \"n_chunk\": n_chunk,\n",
    "            \"meta\": [(uid, arr.shape) for uid, arr in encoded_array_map.items()],\n",
    "        })\n",
    "        del encoded_array_map\n",
    "        with open(out_meta_filepath, \"w\") as f:\n",
    "            json.dump(array_metas, f)\n",
    "        n_chunk += 1\n",
    "        _ = gc.collect()\n",
    "#         break\n",
    "#         if n_chunk == 2:\n",
    "#             break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0135de87",
   "metadata": {},
   "outputs": [],
   "source": [
    "# for 500 files, (max ~120Gb audio in memory)\n",
    "# download ~30s (10 cores), torchload ~45s (100 workers), encode ~80s (8 GPUs)\n",
    "# ~3.5mins for 500 files, ~10 days for 1M hours, $4k on 1x g5.48xlarge\n",
    "create_data(\n",
    "    work_items, \n",
    "    os.path.join(DATA_DIR, \"02_quantization.json\"),\n",
    "    allow_continue=True,\n",
    "    n_gpus=8,\n",
    "    chunksize=500, \n",
    "    gpu_batch_size=16,\n",
    "    n_cores=10,\n",
    "    num_workers_loader=96,\n",
    "    prefetch_factor_loader=2,\n",
    ")\n",
    "# started Mon, Jan 23 at ~2pm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40de5181",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: upload metadata file as well"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6facd9dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !cat tmp/log.txt\n",
    "# !grep -o 's3_filepath' /home/georg/notebooks/web_harvest/harvest/podcasts/data/02_quantization.jsonl | wc -l\n",
    "# !s3cmd ls s3://suno-data/datasets/harvest/podcasts/arrays/ | wc -l"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5f53a9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !s3cmd put /home/georg/notebooks/web_harvest/harvest/podcasts/data/02_quantization.json s3://suno-data/datasets/harvest/podcasts/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b095d6e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54c1da91",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1a9e85c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f16723b0",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "37197edf",
   "metadata": {},
   "source": [
    "## Playground"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "fe01fff7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import contextlib\n",
    "import os\n",
    "import tempfile\n",
    "import threading\n",
    "import time\n",
    "\n",
    "from encodec import EncodecModel\n",
    "from encodec.utils import convert_audio\n",
    "import funcy\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "import torchaudio\n",
    "\n",
    "from suno_utils.audio import Audio\n",
    "from suno_utils.audio.conversion import _convert_audio_ffmpeg\n",
    "\n",
    "SAMPLE_RATE = 24_000\n",
    "\n",
    "MAX_DURATION_S = 1*60*60\n",
    "MAX_ARRAY_LEN = int(MAX_DURATION_S * SAMPLE_RATE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0d10332d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: max_frames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8ebaf39e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_audio(filepath, num_frames=-1):\n",
    "    arr, sr = torchaudio.load(filepath, num_frames=num_frames)\n",
    "    return arr\n",
    "\n",
    "\n",
    "class AudioDataset(Dataset):\n",
    "    \"\"\"Basic loader to allow multi-core loading.\"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        filepaths,\n",
    "        sample_rate=SAMPLE_RATE,\n",
    "        byte_width=2,\n",
    "        n_channels=1,\n",
    "        max_array_len=None,\n",
    "    ):\n",
    "        self.filepaths = filepaths\n",
    "        self.sample_rate = sample_rate\n",
    "        self.byte_width = byte_width\n",
    "        self.n_channels = n_channels\n",
    "        self.max_array_len = max_array_len\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.filepaths)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        if torch.is_tensor(idx):\n",
    "            idx = idx.tolist()\n",
    "        try:\n",
    "            with tempfile.TemporaryDirectory() as tmp_dir:\n",
    "                tmp_wav_filepath = os.path.join(tmp_dir, \"audio.wav\")\n",
    "                _convert_audio_ffmpeg(\n",
    "                    self.filepaths[idx],\n",
    "                    tmp_wav_filepath,\n",
    "                    sample_rate=self.sample_rate,\n",
    "                    byte_width=self.byte_width,\n",
    "                    n_channels=self.n_channels,\n",
    "                )\n",
    "                sample = load_audio(tmp_wav_filepath, num_frames=self.max_array_len)\n",
    "        except:\n",
    "            return None, None\n",
    "        return sample, torch.tensor(sample.shape[-1])\n",
    "\n",
    "\n",
    "def load_audio_mp(filepaths, max_array_len=-1, num_workers=8, prefetch_factor=2):\n",
    "    # https://github.com/pytorch/pytorch/issues/11201\n",
    "    import torch.multiprocessing\n",
    "\n",
    "    torch.multiprocessing.set_sharing_strategy(\"file_system\")\n",
    "    dataset = AudioDataset(filepaths, max_array_len=max_array_len)\n",
    "    dataloader = DataLoader(\n",
    "        dataset,\n",
    "        batch_size=1,\n",
    "        shuffle=False,\n",
    "        num_workers=num_workers,\n",
    "        prefetch_factor=prefetch_factor,\n",
    "    )\n",
    "    arr_list = []\n",
    "    for i_batch, sample_batched in enumerate(dataloader):\n",
    "        arr_list.append(sample_batched[0][0])\n",
    "        # stop after 1 epoch\n",
    "        if len(arr_list) >= len(filepaths):\n",
    "            break\n",
    "    arr_list = arr_list[:len(filepaths)]\n",
    "    return arr_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "0039ea10",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "66.0 s\n"
     ]
    }
   ],
   "source": [
    "t0 = time.time()\n",
    "load_audio_mp(\n",
    "    filepaths, \n",
    "    max_array_len=MAX_ARRAY_LEN, \n",
    "    num_workers=100, \n",
    "    prefetch_factor=2,\n",
    ")\n",
    "t1 = time.time()\n",
    "print(round(t1-t0, 1), \"s\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26075a9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# record\n",
    "# 500 files - 225s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f5912e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import psutil\n",
    ">>> psutil.virtual_memory()\n",
    "svmem(total=16717422592, available=5376126976, percent=67.8, used=10359984128, free=1831890944, active=7191916544, inactive=2325667840, buffers=525037568, cached=4000509952, shared=626225152)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "9760868a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "filepaths = []\n",
    "for n in range(5):\n",
    "    filepaths.append(f\"tmp/{n}.wav\")\n",
    "len(filepaths)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "dd371e3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.tasks.encodec import AudioDataset, DataLoader\n",
    "dataset = AudioDataset(filepaths)\n",
    "dataloader = DataLoader(\n",
    "    dataset,\n",
    "    batch_size=1,\n",
    "    shuffle=False,\n",
    "    num_workers=1,\n",
    "    prefetch_factor=1,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "4a1eeaad",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "1\n",
      "2\n",
      "3\n",
      "4\n"
     ]
    }
   ],
   "source": [
    "n = 0\n",
    "for i_batch, sample_batched in enumerate(dataloader):\n",
    "    print(n)\n",
    "    # stop after 1 epoch\n",
    "#     if len(arr_list) >= len(filepaths):\n",
    "#         break\n",
    "    n += 1\n",
    "    if n == 8:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3f738fd",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3c7220a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "3fec36af",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03fb4f7e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "67230566",
   "metadata": {},
   "outputs": [
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m/tmp/ipykernel_3578550/3921175902.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0msuno_utils\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mutils\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msystem\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mlog_memory\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0mlog_memory\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"log_memory.txt\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0madd_gpu_stats\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrefresh_time_s\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmax_monitor_time_s\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m10\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0;36m60\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0;36m60\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m~/code/glockenspiel/suno_utils/suno_utils/utils/system.py\u001b[0m in \u001b[0;36mlog_memory\u001b[0;34m(filepath, add_gpu_stats, refresh_time_s, max_monitor_time_s)\u001b[0m\n\u001b[1;32m     45\u001b[0m         \u001b[0;32mwith\u001b[0m \u001b[0mopen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfilepath\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"a\"\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     46\u001b[0m             \u001b[0mf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mwrite\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"\\t\"\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mjoin\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mwrite_lines\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;34m\"\\n\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 47\u001b[0;31m         \u001b[0mtime\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msleep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mrefresh_time_s\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     48\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mtime\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtime\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m-\u001b[0m \u001b[0mt0\u001b[0m \u001b[0;34m>\u001b[0m \u001b[0mmax_monitor_time_s\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     49\u001b[0m             \u001b[0;32mbreak\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "from suno_utils.utils.system import log_memory\n",
    "log_memory(\"log_memory.txt\", add_gpu_stats=False, refresh_time_s=2, max_monitor_time_s=10*60*60)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8a77fae1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "8.0K\tlog_memory.txt\r\n"
     ]
    }
   ],
   "source": [
    "!du -hs log_memory.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b96bd8d3",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8177712",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0568e7d6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fa43f39",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60c92d76",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65b18d45",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}