notebooks/mlflow.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a285820-cf70-403a-a39c-0271330526c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ! pip install flowcept[mlflow]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afd1072f-0a6a-4515-b4b9-adc7a5e06cf3",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Set the env var pointing to the conf file where the ports, hostnames, and other conf variables are read from.\n",
    "\n",
    "There is an exemplary conf file available in the `resources` directory in FlowCept repository. You can use it as is if running this Notebook on your local laptop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2cd32bb-8f92-4026-a15b-8cee25189499",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "SQLITE_PATH = 'mlflow.db'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a15ac32f-7bbc-4a14-829c-594467102947",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "## This cell Resets MLFlow Database\n",
    "! rm -f {SQLITE_PATH}\n",
    "! rm -rf mlruns\n",
    "import mlflow\n",
    "mlflow.set_tracking_uri(f\"sqlite:///{SQLITE_PATH}\")\n",
    "mlflow.delete_experiment(mlflow.create_experiment('starter'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ebd2157-dd51-494b-834a-5d6ba2034197",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Sleeps are used because these notebooks are being tested automatically as part of the CI/CD. \n",
    "# In a normal user interaction, these sleeps would not be necessary.\n",
    "from time import sleep"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0290242f-731f-4315-9987-6040c4070687",
   "metadata": {},
   "source": [
    "## Define a simple MLFlow workflow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c08c39cf-8489-47eb-aff3-d96768afc7e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_mlflow_workflow(sqlite_path=SQLITE_PATH, batch_size=64, epochs=10):\n",
    "    import uuid\n",
    "    import mlflow\n",
    "\n",
    "    mlflow.set_tracking_uri(\n",
    "        f\"sqlite:///\" f\"{sqlite_path}\"\n",
    "    )\n",
    "    experiment_name = \"LinearRegression\"\n",
    "    experiment_id = mlflow.create_experiment(\n",
    "        experiment_name + str(uuid.uuid4())\n",
    "    )\n",
    "    with mlflow.start_run(experiment_id=experiment_id) as run:\n",
    "        mlflow.log_params({\"epochs\": epochs})\n",
    "        mlflow.log_params({\"batch_size\": batch_size})\n",
    "        # Actual training code would go here\n",
    "        print(\"Generated training metadata.\")\n",
    "        mlflow.log_metric(\"loss\", 0.04)\n",
    "        return run.info.run_uuid"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b564c09b-7db9-46e3-b079-9799ca860f4f",
   "metadata": {},
   "source": [
    "## Initialize interceptor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7edd571f-9507-496f-8926-6dcc608b2358",
   "metadata": {},
   "outputs": [],
   "source": [
    "from flowcept import MLFlowInterceptor\n",
    "interceptor = MLFlowInterceptor()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c03d7ef6-41cb-4be5-8cda-1aa8ddd33a5a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "assert interceptor.settings.file_path == SQLITE_PATH"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df97ce94-db1a-463e-84fb-66f046ab1716",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Initialize consumer API"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b57ee95c-a061-4f2d-b43a-d7e02d66bda8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from flowcept import Flowcept\n",
    "flowcept = Flowcept(interceptor)\n",
    "flowcept.start()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58ee85e8-c3d1-40ad-95dd-799254d13a93",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Run MLFlow workflow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1c0f6d7-7d47-4c61-9d5b-e5bd398326cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "mlflow_run_id = run_mlflow_workflow(interceptor.settings.file_path, batch_size=18)\n",
    "print(f\"MLflow task id={mlflow_run_id}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1214cbba-ffc4-4b70-8c54-d8560dd7889f",
   "metadata": {},
   "outputs": [],
   "source": [
    "sleep(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a9bf556-8091-4e3a-916f-38a35b097870",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Initialize Query API"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e0dd400-3578-4ace-a3e4-f6feec8b3152",
   "metadata": {},
   "outputs": [],
   "source": [
    "from flowcept import TaskQueryAPI\n",
    "query_api = TaskQueryAPI()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "948e423c-8d1f-4dbc-af7f-16d283e512bd",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Query the task executed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48dbf1ba-f95c-4889-b5fc-1559ed55862f",
   "metadata": {},
   "outputs": [],
   "source": [
    "_filter = {\"task_id\": mlflow_run_id}\n",
    "query_api.query(_filter)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "925781ce-35c0-41e6-9577-46bd2cd0a60f",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Inspect more tasks through the query api"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c90974b-5cb9-4a06-8c4c-608c3d2a4525",
   "metadata": {},
   "outputs": [],
   "source": [
    "from flowcept.commons.utils import get_utc_minutes_ago"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35d40657-94b4-4b34-b82f-e4fb32e823c9",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Get the tasks executed in my experiment in the last 60 minutes\n",
    "\n",
    "This example assumes that you have run the Dask notebook example before. If you haven't run it, just ignore these queries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20bf2264-4f4f-4086-b2ed-7d44e088a3e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "_filter = {\n",
    "    \"utc_timestamp\": { \"$gte\" : get_utc_minutes_ago(60) },\n",
    "    \"campaign_id\": \"super_campaign\"\n",
    "}\n",
    "query_api.query(_filter)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4408e777-8abd-49b8-a9b7-4264eb6004d6",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Get the tasks executed in my experiment in the last 60 minutes that were executed using mlflow and dask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae88d8a0-0c10-4248-9793-daf29c82c488",
   "metadata": {},
   "outputs": [],
   "source": [
    "_filter = {\n",
    "    \"utc_timestamp\": { \"$gte\" : get_utc_minutes_ago(60) },\n",
    "    \"campaign_id\": \"super_campaign\",\n",
    "    \"adapter_id\": {\"$in\": ['mlflow', 'dask']}\n",
    "}\n",
    "docs = query_api.query(_filter)\n",
    "docs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4bfbac82-26e2-4164-987a-a1f2f9abe2ea",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Get the tasks executed in my experiment in the last 60 minutes that generated a batch_size > 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7a854ec-987d-4136-976d-ff144ae6248e",
   "metadata": {},
   "outputs": [],
   "source": [
    "_filter = {\n",
    "    \"utc_timestamp\": { \"$gte\" : get_utc_minutes_ago(120) },\n",
    "    \"campaign_id\": \"super_campaign\",\n",
    "    \"adapter_id\": 'dask',\n",
    "    \"generated.batch_size\": { \"$gte\" : 0 }\n",
    "}\n",
    "docs = query_api.query(filter=_filter)\n",
    "batch_sizes = []\n",
    "for doc in docs:\n",
    "    print(f\"task={doc['task_id']}, generated batch_size={doc['generated']['batch_size']}\")\n",
    "    batch_sizes.append(doc['generated']['batch_size'])\n",
    "batch_sizes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e00a9290-1140-4047-954d-2239d0887a9b",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Now run a new MLFlow task using the batch_sizes generated by the Dask workflow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "072bdd5d-05e5-4387-be0c-e70c902b06d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_sizes = batch_sizes if len(batch_sizes) else [32]  # To use this if you haven't executed the Dask workflow first"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e36f0135-af5e-407e-be5e-66055f260855",
   "metadata": {},
   "outputs": [],
   "source": [
    "for batch_size in batch_sizes:\n",
    "    mlflow_task = run_mlflow_workflow(batch_size=batch_size)\n",
    "    print(mlflow_task)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84c0f0e6-1085-43b1-827c-b26f6dc83aff",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sleep(15)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b56e769-0c27-4e77-9cc3-bbaeb675a575",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Get these tasks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5afab03-897a-4eee-aec5-86743a72d45e",
   "metadata": {},
   "outputs": [],
   "source": [
    "_filter = {\n",
    "    \"task_id\": mlflow_task\n",
    "}\n",
    "docs = query_api.query(filter=_filter)\n",
    "docs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d7d75ed-9dd1-41b3-9276-65d052570771",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert len(docs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02448f6c-3f55-443f-95fd-51ed3c633265",
   "metadata": {},
   "source": [
    "## Stop consumers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fad550a2-ae87-47a7-bd05-50cf6ef38c0b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "flowcept.stop()"
   ]
  }
 ],
 "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.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}