jobflow_to_universal_2.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8676f0ca-5b34-4867-a32d-928f37fb0ed2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from importlib import import_module"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1bd1f75c-a07f-48c9-ae56-22ba674d6d3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jobflow import job, Flow\n",
    "from jobflow.managers.local import run_locally"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "76849db7-1ea1-4d3c-85da-441bb10382f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "@job(data=[\"x\", \"y\", \"z\"])\n",
    "def add_xy(x, y):\n",
    "    return {\"x\": x, \"y\": y, \"z\": x + y}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "446140ce-c375-4cdf-94a1-1ed41d0ee683",
   "metadata": {},
   "outputs": [],
   "source": [
    "@job\n",
    "def add_xyz(x, y, z):\n",
    "    return x + y + z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6d9424db-66a9-43a5-93bd-f28c6499e2d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = 1\n",
    "y = 2\n",
    "z = add_xy(x=x, y=y)\n",
    "w = add_xyz(x=z.output.x, y=z.output.y, z=z.output.z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "0f6fdd39-df5d-4cf8-a987-093a83c4954f",
   "metadata": {},
   "outputs": [],
   "source": [
    "flow = Flow([z, w])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "97b3eb22-cc1d-4462-b94e-18ab7c23ebe3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_function_dict(flow):\n",
    "    return {\n",
    "        job.uuid: job.function\n",
    "        for job in flow.jobs\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "72f39efc-78a0-41f5-8d21-09048dbcd49c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_nodes_dict(function_dict):\n",
    "    nodes_dict, nodes_mapping_dict = {}, {}\n",
    "    for i, [k, v] in enumerate(function_dict.items()):\n",
    "        nodes_dict[i] = v\n",
    "        nodes_mapping_dict[k] = i\n",
    "    \n",
    "    return nodes_dict, nodes_mapping_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "24919dc3-1aa1-4ef8-b54d-96cebf0adcd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_edges_and_extend_nodes(flow_dict, nodes_mapping_dict, nodes_dict):\n",
    "    edges_lst = []\n",
    "    for job in flow_dict['jobs']:\n",
    "        for k, v in job['function_kwargs'].items():\n",
    "            if isinstance(v, dict) and '@module' in v and '@class' in v and '@version' in v:\n",
    "                if len(v['attributes']) == 1:\n",
    "                    edges_lst.append({'target': nodes_mapping_dict[job[\"uuid\"]], 'targetHandle': k, \"source\": nodes_mapping_dict[v['uuid']], 'sourceHandle': v['attributes'][0][1]})\n",
    "                else:\n",
    "                    edges_lst.append({'target': nodes_mapping_dict[job[\"uuid\"]], 'targetHandle': k, \"source\": nodes_mapping_dict[v['uuid']], 'sourceHandle': None})\n",
    "            else:\n",
    "                if v not in nodes_dict.values():\n",
    "                    node_index = len(nodes_dict)\n",
    "                    nodes_dict[node_index] = v\n",
    "                else:\n",
    "                    node_index = {tv: tk for tk, tv in nodes_dict.items()}[v]\n",
    "                edges_lst.append({'target': nodes_mapping_dict[job[\"uuid\"]], 'targetHandle': k, \"source\": node_index, 'sourceHandle': None})\n",
    "    return edges_lst, nodes_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "e5e2e2a2-45e1-4626-b4ab-5bc1ea4ad644",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "([{'target': 0, 'targetHandle': 'x', 'source': 2, 'sourceHandle': None},\n",
       "  {'target': 0, 'targetHandle': 'y', 'source': 3, 'sourceHandle': None},\n",
       "  {'target': 1, 'targetHandle': 'x', 'source': 0, 'sourceHandle': 'x'},\n",
       "  {'target': 1, 'targetHandle': 'y', 'source': 0, 'sourceHandle': 'y'},\n",
       "  {'target': 1, 'targetHandle': 'z', 'source': 0, 'sourceHandle': 'z'}],\n",
       " {0: <function __main__.add_xy(x, y)>,\n",
       "  1: <function __main__.add_xyz(x, y, z)>,\n",
       "  2: 1,\n",
       "  3: 2})"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "flow_dict = flow.as_dict()\n",
    "function_dict = get_function_dict(flow=flow)\n",
    "nodes_dict, nodes_mapping_dict = get_nodes_dict(function_dict=function_dict)\n",
    "edges_lst, nodes_dict = get_edges_and_extend_nodes(flow_dict=flow_dict, nodes_mapping_dict=nodes_mapping_dict, nodes_dict=nodes_dict)\n",
    "edges_lst, nodes_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "c46db577-a0f4-42f0-939d-6a738b7a349c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2025-01-16 18:02:39,727 INFO Started executing jobs locally\n",
      "2025-01-16 18:02:41,075 INFO Starting job - add_xy (3832c38e-98dc-44e7-b6d3-b52a5e86d38c)\n",
      "2025-01-16 18:02:41,076 INFO Finished job - add_xy (3832c38e-98dc-44e7-b6d3-b52a5e86d38c)\n",
      "2025-01-16 18:02:41,076 INFO Starting job - add_xyz (8ec0e25e-93f6-43cf-9f26-36e88184987e)\n",
      "2025-01-16 18:02:41,078 INFO Finished job - add_xyz (8ec0e25e-93f6-43cf-9f26-36e88184987e)\n",
      "2025-01-16 18:02:41,078 INFO Finished executing jobs locally\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'3832c38e-98dc-44e7-b6d3-b52a5e86d38c': {1: Response(output={'x': 1, 'y': 2, 'z': 3}, detour=None, addition=None, replace=None, stored_data=None, stop_children=False, stop_jobflow=False, job_dir=PosixPath('/home/janssen/notebooks/2025/2025-01-16-jobflow-qe'))},\n",
       " '8ec0e25e-93f6-43cf-9f26-36e88184987e': {1: Response(output=6, detour=None, addition=None, replace=None, stored_data=None, stop_children=False, stop_jobflow=False, job_dir=PosixPath('/home/janssen/notebooks/2025/2025-01-16-jobflow-qe'))}}"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result = run_locally(flow)\n",
    "result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "129e94a4-05d1-43f2-8623-2be1e8ed155a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x 1\n",
      "y 2\n",
      "x {'@module': 'jobflow.core.reference', '@class': 'OutputReference', '@version': None, 'uuid': '3832c38e-98dc-44e7-b6d3-b52a5e86d38c', 'attributes': [['a', 'x']], 'output_schema': None}\n",
      "y {'@module': 'jobflow.core.reference', '@class': 'OutputReference', '@version': None, 'uuid': '3832c38e-98dc-44e7-b6d3-b52a5e86d38c', 'attributes': [['a', 'y']], 'output_schema': None}\n",
      "z {'@module': 'jobflow.core.reference', '@class': 'OutputReference', '@version': None, 'uuid': '3832c38e-98dc-44e7-b6d3-b52a5e86d38c', 'attributes': [['a', 'z']], 'output_schema': None}\n"
     ]
    }
   ],
   "source": [
    "for job in flow_dict['jobs']:\n",
    "    for k, v in job['function_kwargs'].items():\n",
    "        print(k, v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dad9d672-a5d7-4acb-bdef-397345d808a4",
   "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.12.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}