{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "a5eba1bd-82df-48c0-8b31-4a4ef33b7e1c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import ray\n",
    "from ray.job_submission import JobSubmissionClient\n",
    "import time\n",
    "import subprocess\n",
    "\n",
    "# Ray cluster information\n",
    "ray_head_ip = \"kuberay-head-svc.kuberay.svc.cluster.local\"\n",
    "ray_head_port = 8265\n",
    "ray_address = f\"http://{ray_head_ip}:{ray_head_port}\"\n",
    "client = JobSubmissionClient(ray_address)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "ad2c8e92-5891-4969-a107-364af5f14d81",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Executors\n",
    "def local_submission(n,m,p):\n",
    "    print(\"Local submission:\")\n",
    "    command = [\"python\", \"ray-matrix_multiplication-example.py\",n,m,p]\n",
    "    subprocess.run(command)\n",
    "\n",
    "def ray_submission(n,m,p,c):\n",
    "    print(\"#################################\\nRay submission:\")\n",
    "    job_id = c.submit_job(\n",
    "        entrypoint=f\"python ray-matrix_multiplication-example.py {n} {m} {p}\",\n",
    "        runtime_env={\n",
    "            \"working_dir\": \"./\", \n",
    "        },\n",
    "        entrypoint_num_cpus=3\n",
    "    )\n",
    "\n",
    "    # Waiting for Ray to finish the job and print the result\n",
    "    while True:\n",
    "        status = c.get_job_status(job_id)\n",
    "        if status in [ray.job_submission.JobStatus.RUNNING, ray.job_submission.JobStatus.PENDING]:\n",
    "            time.sleep(5)\n",
    "        else:\n",
    "            break\n",
    "    try:\n",
    "        logs = c.get_job_logs(job_id) \n",
    "        print(logs)\n",
    "    except RuntimeError as e:\n",
    "        print(f\"Failed to get job logs, please check logs on ray dashboard \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "c00aa8c2-7733-412e-9cc8-cfeccccbbbd8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# matrix1 n x m, matrix2 m x p\n",
    "n = \"10000\"\n",
    "m = \"10000\"\n",
    "p = \"10000\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d637008-8626-40a4-8747-1befbd5d4da5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Local\n",
    "local_submission(n,m,p)\n",
    "\n",
    "# Ray\n",
    "ray_submission(n,m,p,client)\n",
    "ray.shutdown()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d585062f-4c48-4264-8376-7483ad652b00",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Ray",
   "language": "python",
   "name": "ray"
  },
  "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
}