[FEAT] - Raise error when payload is too large and suggest number of …

…partitions (#456)
Nixtla · Sep 3, 2024 · 5bb9e05 · 5bb9e05
1 parent e26e434
commit 5bb9e05
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 0 deletions.
diff --git a/nbs/src/nixtla_client.ipynb b/nbs/src/nixtla_client.ipynb
@@ -702,6 +702,9 @@
     "\n",
     "        ensure_contiguous_arrays(payload)\n",
     "        content = orjson.dumps(payload, option=orjson.OPT_SERIALIZE_NUMPY)\n",
+    "        content_size_mb = len(content) / (1024*1024)\n",
+    "        if content_size_mb > 200:\n",
+    "            raise ValueError(f'The payload is too large. Set num_partitions={math.ceil(content_size_mb / 200)}')\n",
     "        resp = client.post(url=endpoint, content=content)\n",
     "        try:\n",
     "            resp_body = orjson.loads(resp.content)\n",
@@ -2558,6 +2561,21 @@
     ")"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#| hide\n",
+    "# Test large requests raise error and suggest partition number\n",
+    "df = generate_series(20_000, min_length=1_000, max_length=1_000, freq='min')\n",
+    "test_fail(\n",
+    "    lambda: nixtla_client.forecast(df=df, h=1, freq='min', finetune_steps=2),\n",
+    "    contains=\"num_partitions\"\n",
+    ")"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},

diff --git a/nixtla/nixtla_client.py b/nixtla/nixtla_client.py
@@ -637,6 +637,11 @@ def ensure_contiguous_arrays(d: Dict[str, Any]) -> None:
 
         ensure_contiguous_arrays(payload)
         content = orjson.dumps(payload, option=orjson.OPT_SERIALIZE_NUMPY)
+        content_size_mb = len(content) / (1024 * 1024)
+        if content_size_mb > 200:
+            raise ValueError(
+                f"The payload is too large. Set num_partitions={math.ceil(content_size_mb / 200)}"
+            )
         resp = client.post(url=endpoint, content=content)
         try:
             resp_body = orjson.loads(resp.content)