Tomography-flow.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0d3813a6",
   "metadata": {},
   "source": [
    "<img src=\"img/automation_using_flows_header.png\">\n",
    "\n",
    "# Cross-Facility Tomography Reconstructions\n",
    "<img src=\"img/lightsources.png\">\n",
    "\n",
    "This notebook shows how [Tomopy](https://tomopy.readthedocs.io/en/latest/) reconstructions can be performed at different facilities using [Globus Flows](https://www.globus.org/globus-flows-service). "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ca5d8f2",
   "metadata": {},
   "source": [
    "# Using the Flows GUI"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d216f3f",
   "metadata": {},
   "source": [
    "The Globus Flows web app can generate a `start` page to easily invoke flows. To create this page you need to attach an input schema to the flow. This is shown at the end of the notebook.\n",
    "\n",
    "<img src=\"img/input1.png\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07d5e761",
   "metadata": {},
   "source": [
    "### Monitoring the flow\n",
    "Once the flow has been started you can monitor the run. The events tab shows each step of the flow and provides details regarding the input and output of each action.\n",
    "\n",
    "Here you can see the input to the Tomopy step of the flow, showing the input values that are passed into the function to execute.\n",
    "\n",
    "<img src=\"img/running.png\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fcb8ae33",
   "metadata": {},
   "source": [
    "# A Tomopy Flow\n",
    "\n",
    "We have defined a Globus Flow to run Tomopy. The flow definition can be [viewed in the Flows IDE here.](https://globus.github.io/flows-ide?d=N4Igwg9gtlCmB2AXEAuEAVATgQ3gZwDNZMACXAExMwFd4T1oIAHATxABoQBlRbTRAILI0AMWwBjAJYAbSYhZgAFhEnjYHbr0Sw8qUGKmz5SlWr0YWTdWhOr1nW2t0oA2qABqfSdgBG06yAAJAB0kvBM1IjB4tAR2gD6eHL2mphhAOYAogCO1NjSziACADJgIhoAcrAAHsJFpSKQ8ASS6SAAvgC6nAAisATY1NJ1FZkASlxgTS1t7ZwlZdOt5uiWAQAK2Hi6nJs4cNqYzqBJ2sGBqEGh4ZHRsZGw8THNrXjBC+VzIGM6Q4ibiEUlxCYTidygcRSVVqlywuEIxAAkjdkF9RhMphAXm0UKBVlZLpttho9tgDsRjiBTrBzsDrmCYhCHk8sTM3ujJh1OD88H8AUC0CCUeDIZUanU4fgiJhkXEuRgcFKkSjzJAYAgJYqEaRQZESC1-Bp8QEBOJEJIsRpTeasQBVTDSS6KRCIJh4FAAeg9Eht+GCg0Q0GwZ3S0ggPmobwgmHSHsQWulcYTxBJfDJsEOlLwEGomDU8QQ5CYKiQ8Uk5Fpgvpt2zubUoXIGnIOnN8GDFvgBfgRZLiDLFYuVd1UUZkOC1OC8fh0q7PbCyE4U6VmDL2igzjcVJzeceTGDimCAF5LgADIVg2s74J7wEnkgAahIZ+rUWH8QNCHTJ6bLbC7ax8Q3gex5oM+w4ig847JOOvDpBkgH7nej5gcKb4fm2cDfpwmCwOIuZJAAbtYAwFLAXRfDyfL7nSkrarKkQaAA6tgcjoJIcCoAAzAADNxnDQhKjCsPKDBQMwbC4uA0BwEghLEAQ0ZQPQQksGQbbSCwSQ7BYBJoNaHZWmaHb2o6aDOq67peqOkE+h2byhuGkbBNGsYENU2BMKm+wZhS5iFsW86Vlc4HWWcE6hY8-m9hoBC0EZWJBeetwRe+cW+v2GgANYAO58OkWbYER74QNIzaYIlL4QWF0E5q6kQIYCGgxNI-jxZ2QEVSF9zVWczWtelQExTIn5wEep5JSO3U0uFJX9R2DWKEhT4TVV03QXgsHwUBS0oWCaHDRhsBYSAjJQBQnXCilsXwG18Q5XlbynRQGg4c8G00Ld8hWBdDJTalN3pfdMZvK9WLvdQn1rB0FG-MM-I0SpADykRypwzGsex1gAGy8fx4qwsmmCUcMuhfLR0rE4glJqjJmrTsQVCw1TRpQ3pbWGb6JlOi6bqet6bVvAGQYhmGEZRjGSb05gkvLl56aZuYl75lF879j9yVTVBZxLtqs4BbJnDNhtf7pSrpblurURKzS5YaDrM7JOuqCbtbC0VaJ4nI3VUTNrwMgg0zLjcZ0LgAIydD+xttgN+5jaBK04bywwLTtIQe6wXtgr7LEFMEid-EHIfB8dr34ZIRGoCReBkZ0MNJ-81FVuTxAiMN2no4gbEcSgPF8SAmTdqg8bUGR7TtEAA)\n",
    "\n",
    "This flow starts by selecting the facility to perform the tasks. The decision is based on an input to the flow (either `ALCF` or `NERSC`). The appropriate state is then injected into the subsequent steps to direct where data are transferred and specify the Compute endpoint to use.\n",
    "\n",
    "<img src=\"img/flow_definition.png\" style=\"width: 300px;\"/>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3989f718",
   "metadata": {},
   "source": [
    "Here we create a Globus FlowsClient to securely interact with the Flows service. This will prompt a login and paste a token into the notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a29de715-5a27-4671-aaa2-0809f3fbf7ca",
   "metadata": {
    "metadata": {}
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import time\n",
    "import globus_sdk\n",
    "\n",
    "from globus_sdk.experimental.globus_app import UserApp\n",
    "from globus_sdk.scopes import GCSCollectionScopeBuilder\n",
    "\n",
    "# Tutorial client ID. We recommend replacing this with your own client for any production use-cases\n",
    "# Create your own at developers.globus.org\n",
    "CLIENT_ID = \"7ca21f4a-11de-4d97-8f84-cb66f7459981\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f72abbf6",
   "metadata": {},
   "outputs": [],
   "source": [
    "my_app = UserApp(\"tomography-user-app\", client_id=CLIENT_ID)\n",
    "\n",
    "flows_client = globus_sdk.FlowsClient(app=my_app)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "862f12d2-0a70-4ed4-bae2-657256cdf4e2",
   "metadata": {},
   "source": [
    "Specify the ID of the previously registered flow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cfb99f8e-8147-45ff-968d-2465f444dac6",
   "metadata": {},
   "outputs": [],
   "source": [
    "flow_id = '081f12ba-cda1-48ce-89b5-ae3dd68a5689'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72ed205b",
   "metadata": {},
   "source": [
    "## Creating a Tomopy Function\n",
    "\n",
    "To run Tomopy in the flow we need to define a Tomopy function and register it with Globus Compute.\n",
    "\n",
    "**Note, this uses my personal compute endpoint, not a service account endpoint. Access to this endpoint is restricted.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d28edc81",
   "metadata": {},
   "outputs": [],
   "source": [
    "import globus_compute_sdk\n",
    "\n",
    "compute_ep = '70f76338-f3a2-4992-9c44-f38af073bb50'#'0f44e1f7-baa5-4dfd-a2f6-73051004cfaa'\n",
    "gce = globus_compute_sdk.Executor(endpoint_id=compute_ep)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c3de22e",
   "metadata": {},
   "source": [
    "Define a function to call tomocupy.\n",
    "\n",
    "This is a simple wrapper function that uses subprocess to invoke a templated bash command to run the tomopy cli."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aeabdb3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def tomopy_wrapper(filename, command=\"recon\", reconstruction_type=\"slice\", \n",
    "                   save_folder=\"_rec\", collection_path=\"/eagle\",\n",
    "                  ):\n",
    "    import subprocess\n",
    "    import glob\n",
    "    import os\n",
    "    \n",
    "    cmd = f\"tomopy {command} --file-name {filename} --reconstruction-type {reconstruction_type} --save-folder={save_folder}\"\n",
    "    res = subprocess.run(cmd.split(\" \"), stdout=subprocess.PIPE, stderr=subprocess.PIPE)\n",
    "    \n",
    "    # Get the most recent reconstruction file\n",
    "    list_of_files = glob.glob(f'{save_folder}/slice_rec/*')\n",
    "    recon_file = max(list_of_files, key=os.path.getctime)\n",
    "\n",
    "    recon_filename =  os.path.basename(recon_file)\n",
    "    \n",
    "    # The path from a Globus endpoint, this is the output path without the collection prefix\n",
    "    transfer_path = recon_file.replace(collection_path, \"\")\n",
    "    \n",
    "    return recon_filename, transfer_path,  res.returncode, res.stdout.decode(\"utf-8\"), res.stderr.decode(\"utf-8\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c54ef7b",
   "metadata": {},
   "source": [
    "Register the function with Compute. This returns a function ID that can be passed to the flow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9c249df",
   "metadata": {},
   "outputs": [],
   "source": [
    "tomo_function_id = gce.client.register_function(tomopy_wrapper)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2843c862-792e-4143-a7fa-742627470a9d",
   "metadata": {},
   "source": [
    "### Test the function\n",
    "\n",
    "Test the function by running directly with Globus Compute. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39e36038",
   "metadata": {},
   "outputs": [],
   "source": [
    "filename = \"/global/homes/r/rchard/IRI/test_data_1.h5\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4436ee9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "future = gce.submit(tomopy_wrapper, filename, save_folder=\"/global/homes/r/rchard/IRI/outputs\")\n",
    "future.result()[1]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11b217f9",
   "metadata": {},
   "source": [
    "# Define Facility information\n",
    "\n",
    "Switching between facilities relies on inputting:\n",
    "- Globus Transfer endpoints for source and destination\n",
    "- The Globus Compute endpoint deployed at the target facility\n",
    "- Path information relating to the facility, such as project space and the mapping to the Globus Collection\n",
    "\n",
    "Here we define configs to capture this information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab738aef-c99d-4c50-9312-5850a16daccf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configs for Compute resources\n",
    "compute_configs = {\n",
    "    \"ALCF\": {\n",
    "        \"compute_endpoint\": \"0f44e1f7-baa5-4dfd-a2f6-73051004cfaa\",\n",
    "        \"transfer_endpoint\": \"05d2c76a-e867-4f67-aa57-76edeb0beda0\", # dtn#eagle\n",
    "        \"collection_path\": \"/eagle\",\n",
    "        \"staging_path\": \"/APSDataAnalysis/rchard/IRI/\",\n",
    "        \"output_path\": \"/eagle/APSDataAnalysis/rchard/IRI/outputs/\",\n",
    "    },\n",
    "    \"NERSC\": {\n",
    "        \"compute_endpoint\": \"70f76338-f3a2-4992-9c44-f38af073bb50\",\n",
    "        \"transfer_endpoint\": \"6bdc7956-fc0f-4ad2-989c-7aa5ee643a79\",\n",
    "        \"collection_path\": \"\",\n",
    "        \"staging_path\": \"/global/homes/r/rchard/IRI/\",\n",
    "        \"output_path\": \"/global/homes/r/rchard/IRI/outputs/\",\n",
    "    },\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c95f3797-6bd1-474b-8545-c1d292010bda",
   "metadata": {},
   "source": [
    "## Running the flow\n",
    "\n",
    "Define the flow input. This describes the file to act on, which site to use, site configs, and the reconstruction parameters. The flow will use the `compute_site` to inject the associated site configuration (e.g., endpoints) into the transfer and compute steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "282b155f",
   "metadata": {},
   "outputs": [],
   "source": [
    "input_filename = \"test_data_1.h5\"\n",
    "compute_site = \"NERSC\"\n",
    "\n",
    "flow_input = {\n",
    "    \"input\": {\n",
    "        \"input_filename\": input_filename,\n",
    "        \"compute_site\": compute_site,\n",
    "        \"compute_configs\": compute_configs,\n",
    "        \"source\": {\n",
    "            \"id\": \"a17d7fac-ce06-4ede-8318-ad8dc98edd69\",\n",
    "            \"path\": \"/Tomography/inputs/\"\n",
    "        },\n",
    "        \"result_path\": \"/Tomography/outputs/\",\n",
    "        \"compute_function_id\": tomo_function_id,\n",
    "        \"compute_function_kwargs\": {\n",
    "            \"command\": \"recon\",\n",
    "            \"reconstruction_type\": \"slice\",\n",
    "        }\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38348061-2ce1-4dde-8775-79ebdc89cc09",
   "metadata": {},
   "source": [
    "Create a SpecificFlowClient to interact with the flow. This is used to start and monitor the run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c66af4a9-a131-4131-a6b3-5cd21ad5b04e",
   "metadata": {},
   "outputs": [],
   "source": [
    "specific_flow_client = globus_sdk.SpecificFlowClient(\n",
    "    flow_id=flow_id,\n",
    "    app=my_app,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4fa733e0-ccad-41de-abe5-f81390b91e65",
   "metadata": {},
   "source": [
    "Add scopes to the flow to use Mapped Collections if Guest Collections are not available. Without this step, the flow will pause before performing the transfer for you to review consents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ffa89e6-93c6-4d37-bb47-21c05f87250c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the data access scope Mapped collections\n",
    "dest_transfer_endpoint = flow_input['input']['compute_configs'][compute_site]['transfer_endpoint']\n",
    "dest_access_scope = GCSCollectionScopeBuilder(dest_transfer_endpoint).make_mutable(\"data_access\", optional=True)\n",
    "\n",
    "transfer_scope = globus_sdk.scopes.TransferScopes.make_mutable(\"all\")\n",
    "transfer_scope.add_dependency(dest_access_scope)\n",
    "\n",
    "# Add the data access scopes as dependencies to the flow\n",
    "flow_scope = specific_flow_client.scopes.make_mutable(\"user\")\n",
    "flow_scope.add_dependency(transfer_scope)\n",
    "\n",
    "my_app.add_scope_requirements({'flow': [flow_scope]})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca56ad14-ecb3-45f9-8952-fc7213619aac",
   "metadata": {},
   "source": [
    "Start the flow. Include a label and tags to better manage and filter runs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10f8aaa5-cc35-4abb-8bb5-590e29069154",
   "metadata": {},
   "outputs": [],
   "source": [
    "run = specific_flow_client.run_flow(\n",
    "  body=flow_input,\n",
    "  label=\"Tomopy run\",\n",
    "  tags=['Tomopy', 'IRI', 'example', compute_site]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6b72db7",
   "metadata": {},
   "outputs": [],
   "source": [
    "run_id = run['run_id']\n",
    "run_status = run['status']\n",
    "print(\"This flow can be monitored in the Web App:\")\n",
    "print(f\"https://app.globus.org/runs/{run_id}\")\n",
    "print(f\"Flow run started with ID: {run_id} - Status: {run_status}\")\n",
    "\n",
    "# Poll the Flows service to check on the status of the run\n",
    "while run_status == 'ACTIVE':\n",
    "    time.sleep(5)\n",
    "    run = flows_client.get_run(run_id)\n",
    "    run_status = run['status']\n",
    "    print(f'Run status: {run_status}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3980bba",
   "metadata": {},
   "outputs": [],
   "source": [
    "run['details']['output']['TomopyOutput']['details']['result'][0][0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4557ffc7",
   "metadata": {},
   "source": [
    "## Attaching an input schema\n",
    "\n",
    "We can use a JSON input schema to both generate the Web interface and provide additional handrails when starting the flow. Here we define the schema and update the flow to include it.\n",
    "\n",
    "Example schema's can be found here: https://github.com/globus/globus-flows-trigger-examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ccad565",
   "metadata": {},
   "outputs": [],
   "source": [
    "schema = {\n",
    "    \"required\": [\n",
    "        \"input\"\n",
    "    ],\n",
    "    \"properties\": {\n",
    "        \"input\": {\n",
    "            \"type\": \"object\",\n",
    "            \"required\": [\n",
    "                \"source\",\n",
    "                \"input_filename\",\n",
    "                \"result_path\",\n",
    "                \"compute_site\",\n",
    "                \"compute_configs\",\n",
    "                \"compute_function_id\",\n",
    "                \"compute_function_kwargs\"\n",
    "            ],\n",
    "            \"properties\": {\n",
    "                \"source\": {\n",
    "                    \"type\": \"object\",\n",
    "                    \"title\": \"Select source collection and path\",\n",
    "                    \"description\": \"The source collection and path (path MUST end with a slash)\",\n",
    "                    \"format\": \"globus-collection\",\n",
    "                    \"required\": [\n",
    "                        \"id\",\n",
    "                        \"path\"\n",
    "                    ],\n",
    "                    \"properties\": {\n",
    "                        \"id\": {\n",
    "                            \"type\": \"string\",\n",
    "                            \"format\": \"uuid\"\n",
    "                        },\n",
    "                        \"path\": {\n",
    "                            \"type\": \"string\"\n",
    "                        }\n",
    "                    },\n",
    "                    \"additionalProperties\": False\n",
    "                },\n",
    "                \"input_filename\": {\n",
    "                    \"type\" : \"string\",\n",
    "                    \"description\": \"The file to process\",\n",
    "                    \"default\": \"test_data_1.h5\",\n",
    "                },\n",
    "                \"result_path\": {\n",
    "                    \"type\": \"string\",\n",
    "                    \"description\": \"Path to return results at source endpoint\",\n",
    "                    \"default\": \"/Tomography/outputs/\"\n",
    "                },\n",
    "                \"compute_site\": {\n",
    "                    \"type\" : \"string\",\n",
    "                    \"description\": \"Facility: ALCF, NERSC\",\n",
    "                    \"default\": \"ALCF\",\n",
    "                    \"enum\" : [\n",
    "                        \"ALCF\",\n",
    "                        \"NERSC\"\n",
    "                    ]\n",
    "                },\n",
    "                \"compute_configs\": {\n",
    "                    \"type\": \"object\",\n",
    "                    \"description\": \"Configs for various Compute sites\",\n",
    "                    \"default\": compute_configs\n",
    "                },\n",
    "                \"compute_function_id\": {\n",
    "                    \"type\": \"string\",\n",
    "                    \"format\": \"uuid\",\n",
    "                    \"description\": \"Tomopy function UUID\",\n",
    "                    \"default\": tomo_function_id\n",
    "                },\n",
    "                \"compute_function_kwargs\": {\n",
    "                    \"type\": \"object\",\n",
    "                    \"title\": \"Function Inputs\",\n",
    "                    \"description\": \"Inputs to pass to Tomopy\",\n",
    "                    \"required\": [\n",
    "                        \"command\",\n",
    "                        \"reconstruction_type\",\n",
    "                    ],\n",
    "                    \"properties\": {\n",
    "                        \"command\": {\n",
    "                            \"type\" : \"string\",\n",
    "                            \"description\": \"Reconstruction command: recon, recon_steps\",\n",
    "                            \"default\": \"recon\",\n",
    "                            \"enum\" : [\n",
    "                                \"recon\",\n",
    "                                \"recon_steps\"\n",
    "                            ]\n",
    "                        },\n",
    "                        \"reconstruction_type\": {\n",
    "                            \"type\": \"string\",\n",
    "                            \"description\": \"Reconstruction type: slice, full\",\n",
    "                            \"default\": \"slice\",\n",
    "                            \"enum\": [\n",
    "                                \"slice\", \n",
    "                                \"full\"\n",
    "                            ]\n",
    "                        },\n",
    "                    },\n",
    "                    \"additionalProperties\": False\n",
    "                }\n",
    "            },\n",
    "            \"additionalProperties\": False\n",
    "        },    \n",
    "    },\n",
    "    \"additionalProperties\": False\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2e2b247",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "flows_client.update_flow(flow_id, input_schema=schema)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79aec70b",
   "metadata": {},
   "source": [
    "## Try it!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3dd4460",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f'https://app.globus.org/flows/{flow_id}')"
   ]
  }
 ],
 "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}