diff --git a/other/materials_designer/workflows/Introduction.ipynb b/other/materials_designer/workflows/Introduction.ipynb index 51bef958b..5a9cd46d3 100644 --- a/other/materials_designer/workflows/Introduction.ipynb +++ b/other/materials_designer/workflows/Introduction.ipynb @@ -105,7 +105,8 @@ "## 9. Custom\n", "\n", "### 9.1. Python / Shell\n", - "#### 9.1.1. Custom Python and Shell workflows. *(to be added)*\n" + "#### [9.1.1. Custom Python calculation.](custom_python_calculation.ipynb)\n", + "#### [9.1.2. Custom Shell calculation.](custom_shell_calculation.ipynb)\n" ] }, { diff --git a/other/materials_designer/workflows/custom_python_calculation.ipynb b/other/materials_designer/workflows/custom_python_calculation.ipynb new file mode 100644 index 000000000..bd63fe237 --- /dev/null +++ b/other/materials_designer/workflows/custom_python_calculation.ipynb @@ -0,0 +1,614 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": {}, + "source": [ + "# Custom Python Calculation\n", + "\n", + "Run your own Python script against one or more materials on the Mat3ra platform. The script and any\n", + "files it needs are uploaded to your object storage folder, a workflow fetches them onto the compute\n", + "node alongside the material, and whatever the script prints comes back as the result.\n", + "\n", + "

Usage

\n", + "\n", + "1. Set the materials and, if the script needs them, its dependencies and data files in cell 1.2.\n", + " below.\n", + "1. Replace the script in cell 1.3. with your own (or keep the default one).\n", + "1. Click \"Run\" > \"Run All\" to run all cells.\n", + "1. Wait for the jobs to complete.\n", + "1. Scroll down to view the results.\n", + "\n", + "## Summary\n", + "\n", + "1. Set up the environment and parameters: install packages (JupyterLite only) and configure the\n", + " materials, the script, its dependencies, its data files, compute resources and job.\n", + "1. Authenticate and initialize API client: authenticate via browser, initialize the client, then\n", + " select account and project.\n", + "1. Create materials: materials are read from the `../uploads` folder — place files there manually or\n", + " run a material creation notebook first. If a material is not found by name, Standata is used as a\n", + " fallback. Each material is then saved to the platform.\n", + "1. Create workflow: upload the script and its data files, then assemble a workflow that fetches\n", + " them, fetches the material, and runs the script. Optionally save the workflow to the collection.\n", + "1. Configure compute: get list of clusters and create compute configuration with selected cluster,\n", + " queue, and number of processors.\n", + "1. Create one job per material from the material, workflow, project and compute configuration.\n", + "1. Submit the jobs and monitor the status: submit and wait for completion.\n", + "1. Retrieve results: read each job's standard output and display the values it printed.\n", + "\n", + "## How the script receives its inputs\n", + "\n", + "Everything lands in the job's working directory, so the script reads it all by **relative path**:\n", + "\n", + "| File | Written by | Contents |\n", + "| --- | --- | --- |\n", + "| `material.json` | the workflow | the job's material, as stored on the platform |\n", + "| your data files | this notebook | uploaded verbatim from `USER_ASSET_FILES` |\n", + "\n", + "The script's standard output is the result. Print JSON and this notebook renders it as a table." + ] + }, + { + "cell_type": "markdown", + "id": "1", + "metadata": {}, + "source": [ + "## 1. Set up the environment and parameters\n", + "### 1.1. Install packages (JupyterLite)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "from mat3ra.notebooks_utils.packages import install_packages\n", + "\n", + "await install_packages(\"made|api_examples\")" + ] + }, + { + "cell_type": "markdown", + "id": "3", + "metadata": {}, + "source": [ + "### 1.2. Set parameters and configurations for the workflow and job\n", + "\n", + "`USER_ASSET_FILES` names data files the script opens, of any type. Put them in the `../uploads`\n", + "folder first — drag them into the JupyterLite file browser — and this notebook uploads them\n", + "alongside the script." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "from datetime import datetime\n", + "\n", + "from mat3ra.ide.compute import QueueName\n", + "\n", + "# 2. Auth and organization parameters\n", + "# Set organization name to use it as the owner, otherwise your personal account is used\n", + "ORGANIZATION_NAME = None\n", + "\n", + "# 3. Material parameters\n", + "FOLDER = \"../uploads\"\n", + "MATERIAL_NAMES = [\"Silicon\"] # One job is created per material\n", + "\n", + "# 4. Script parameters\n", + "USER_REQUIREMENTS = [] # e.g. [\"numpy<2\"], installed into a virtual environment on the compute node\n", + "USER_ASSET_FILES = [] # data files the script opens, taken from FOLDER\n", + "\n", + "# 5. Workflow parameters\n", + "WORKFLOW_SEARCH_TERM = \"custom_script.json\" # Search term for Workflows Standata\n", + "APPLICATION_NAME = \"python\"\n", + "MY_WORKFLOW_NAME = \"Custom Python Calculation\"\n", + "save_to_collection = True\n", + "\n", + "# 6. Compute parameters\n", + "CLUSTER_NAME = None # specify full or partial name i.e. \"cluster-001\" to select\n", + "QUEUE_NAME = QueueName.D\n", + "PPN = 1\n", + "\n", + "# 7. Job parameters\n", + "timestamp = datetime.now().strftime(\"%Y-%m-%d %H:%M\")\n", + "POLL_INTERVAL = 30 # seconds" + ] + }, + { + "cell_type": "markdown", + "id": "5", + "metadata": {}, + "source": [ + "### 1.3. Set the script to run\n", + "\n", + "This is the calculation. It runs on the compute node with `material.json` and any data files beside\n", + "it, and whatever it prints becomes the result. Replace it with your own.\n", + "\n", + "The script is uploaded as a file and fetched onto the node, never inlined into the workflow, so its\n", + "contents reach Python exactly as written - text that looks like a template placeholder is left\n", + "alone." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "USER_SCRIPT = r\"\"\"\n", + "import json\n", + "\n", + "material = json.load(open(\"material.json\"))\n", + "elements = [element[\"value\"] for element in material[\"basis\"][\"elements\"]]\n", + "\n", + "print(json.dumps({\"n_atoms\": len(elements), \"elements\": sorted(set(elements))}))\n", + "\"\"\"" + ] + }, + { + "cell_type": "markdown", + "id": "7", + "metadata": {}, + "source": [ + "## 2. Authenticate and initialize API client\n", + "### 2.1. Authenticate\n", + "Authenticate in the browser and have credentials stored in environment variable \"OIDC_ACCESS_TOKEN\"." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "metadata": {}, + "outputs": [], + "source": [ + "from mat3ra.notebooks_utils.auth import authenticate\n", + "\n", + "\n", + "await authenticate()" + ] + }, + { + "cell_type": "markdown", + "id": "9", + "metadata": {}, + "source": [ + "### 2.2. Initialize API Client" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10", + "metadata": {}, + "outputs": [], + "source": [ + "from mat3ra.api_client import APIClient\n", + "\n", + "client = APIClient.authenticate()\n", + "client" + ] + }, + { + "cell_type": "markdown", + "id": "11", + "metadata": {}, + "source": [ + "### 2.3. Select account to work under" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12", + "metadata": {}, + "outputs": [], + "source": [ + "client.list_accounts()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "13", + "metadata": {}, + "outputs": [], + "source": [ + "selected_account = client.my_account\n", + "\n", + "if ORGANIZATION_NAME:\n", + " selected_account = client.get_account(name=ORGANIZATION_NAME)\n", + "\n", + "ACCOUNT_ID = selected_account.id\n", + "print(f\"Using account: {selected_account.name} ({ACCOUNT_ID})\")" + ] + }, + { + "cell_type": "markdown", + "id": "14", + "metadata": {}, + "source": [ + "### 2.4. Select project" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "15", + "metadata": {}, + "outputs": [], + "source": [ + "projects = client.projects.list({\"isDefault\": True, \"owner._id\": ACCOUNT_ID})\n", + "project_id = projects[0][\"_id\"]\n", + "print(f\"✅ Using project: {projects[0]['name']} ({project_id})\")" + ] + }, + { + "cell_type": "markdown", + "id": "16", + "metadata": {}, + "source": [ + "## 3. Create materials\n", + "### 3.1. Load materials from local files (or Standata)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "17", + "metadata": {}, + "outputs": [], + "source": [ + "from mat3ra.made.material import Material\n", + "from mat3ra.standata.materials import Materials\n", + "from mat3ra.notebooks_utils.ipython.entity.material.visualize import visualize_materials as visualize\n", + "from mat3ra.notebooks_utils.material import load_material_from_folder\n", + "\n", + "materials = [\n", + " load_material_from_folder(FOLDER, name) or Material.create(Materials.get_by_name_first_match(name))\n", + " for name in MATERIAL_NAMES\n", + "]\n", + "visualize([{\"material\": material, \"title\": material.name} for material in materials])" + ] + }, + { + "cell_type": "markdown", + "id": "18", + "metadata": {}, + "source": [ + "### 3.2. Save materials to the platform" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "19", + "metadata": {}, + "outputs": [], + "source": [ + "from mat3ra.notebooks_utils.core.entity.material.api import get_or_create_material\n", + "\n", + "saved_materials = [\n", + " Material.create(get_or_create_material(client, material, ACCOUNT_ID)) for material in materials\n", + "]" + ] + }, + { + "cell_type": "markdown", + "id": "20", + "metadata": {}, + "source": [ + "### 4.1. Upload the script and its data files\n", + "\n", + "The files go to your account's object storage folder (\"Dropbox\"), which the compute node reads them\n", + "from. Text travels inside the request; anything else — a model checkpoint, an archive, a large file\n", + "— is sent to a URL the platform signs for it, straight to storage, so there is no size or format\n", + "limit beyond what the browser can hold in memory." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "21", + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "\n", + "from mat3ra.notebooks_utils.core.entity.file.api import upload_files\n", + "\n", + "files_to_upload = {\"user_script.py\": USER_SCRIPT}\n", + "for name in USER_ASSET_FILES:\n", + " if name in files_to_upload:\n", + " raise ValueError(f\"Rename '{name}' in USER_ASSET_FILES: this notebook already uploads a file by that name.\")\n", + " path = os.path.join(FOLDER, name)\n", + " if not os.path.exists(path):\n", + " raise FileNotFoundError(f\"'{name}' is listed in USER_ASSET_FILES but is not in {FOLDER}.\")\n", + " with open(path, \"rb\") as file:\n", + " content = file.read()\n", + " try:\n", + " files_to_upload[name] = content.decode(\"utf-8\") # text travels in the request body\n", + " except UnicodeDecodeError:\n", + " files_to_upload[name] = content # anything else goes to storage through a signed URL\n", + "\n", + "uploaded_files = upload_files(client, files_to_upload, ACCOUNT_ID)" + ] + }, + { + "cell_type": "markdown", + "id": "22", + "metadata": {}, + "source": [ + "### 4.2. Create workflow from standard workflows and preview it\n", + "\n", + "The `Custom Python Script` workflow already carries the unit chain this needs: fetch the uploaded\n", + "files, fetch the material, run the script. Three things are filled in per job — the objects to\n", + "fetch, the runner that hands the script its inputs, and the dependency list." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "23", + "metadata": {}, + "outputs": [], + "source": [ + "from mat3ra.standata.applications import ApplicationStandata\n", + "from mat3ra.ade.application import Application\n", + "from mat3ra.standata.workflows import WorkflowStandata\n", + "from mat3ra.wode.workflows import Workflow\n", + "from mat3ra.notebooks_utils.core.entity.file.api import to_object_storage_input\n", + "from mat3ra.notebooks_utils.core.entity.workflow.api import CUSTOM_SCRIPT_RUNNER, set_execution_unit_input\n", + "from mat3ra.notebooks_utils.ipython.entity.workflow.visualize import visualize_workflow\n", + "\n", + "app_config = ApplicationStandata.get_by_name_first_match(APPLICATION_NAME)\n", + "app = Application(**app_config)\n", + "\n", + "workflow_config = WorkflowStandata.filter_by_application(app.name).get_by_name_first_match(WORKFLOW_SEARCH_TERM)\n", + "workflow_config[\"name\"] = MY_WORKFLOW_NAME\n", + "subworkflow = workflow_config[\"subworkflows\"][0]\n", + "subworkflow[\"name\"] = MY_WORKFLOW_NAME\n", + "units = {unit[\"name\"]: unit for unit in subworkflow[\"units\"]}\n", + "\n", + "units[\"io-user-files\"][\"input\"] = [to_object_storage_input(file) for file in uploaded_files]\n", + "\n", + "set_execution_unit_input(units[\"custom_script\"], \"script.py\", CUSTOM_SCRIPT_RUNNER)\n", + "set_execution_unit_input(units[\"custom_script\"], \"requirements.txt\", \"\\n\".join(USER_REQUIREMENTS) + \"\\n\")\n", + "\n", + "workflow = Workflow.create(workflow_config)\n", + "visualize_workflow(workflow)" + ] + }, + { + "cell_type": "markdown", + "id": "24", + "metadata": {}, + "source": [ + "### 4.3. Save workflow to collection\n", + "\n", + "Saving it makes the workflow reusable: it stays in your collection pointing at the uploaded files,\n", + "so it can be run again — from the UI or another notebook — against any material. Re-running this\n", + "notebook overwrites the uploaded files in place, so a saved workflow picks up an edited script\n", + "without being rebuilt." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "25", + "metadata": {}, + "outputs": [], + "source": [ + "saved_workflow = None\n", + "if save_to_collection:\n", + " saved_workflow = Workflow.create(\n", + " client.workflows.create(workflow.to_dict_without_special_keys(), owner_id=ACCOUNT_ID)\n", + " )\n", + " print(f\"✅ Workflow saved to collection: {saved_workflow.id}\")" + ] + }, + { + "cell_type": "markdown", + "id": "26", + "metadata": {}, + "source": [ + "## 5. Create the compute configuration\n", + "### 5.1. Get list of clusters" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "27", + "metadata": {}, + "outputs": [], + "source": [ + "clusters = client.clusters.list()\n", + "print(f\"Available clusters: {[c['hostname'] for c in clusters]}\")" + ] + }, + { + "cell_type": "markdown", + "id": "28", + "metadata": {}, + "source": [ + "### 5.2. Create compute configuration for the jobs" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "29", + "metadata": {}, + "outputs": [], + "source": [ + "from mat3ra.ide.compute import Compute\n", + "\n", + "# Select cluster: use specified name if provided, otherwise use first available\n", + "if CLUSTER_NAME:\n", + " cluster = next((c for c in clusters if CLUSTER_NAME in c[\"hostname\"]), None)\n", + "else:\n", + " cluster = clusters[0]\n", + "\n", + "compute = Compute(\n", + " cluster=cluster,\n", + " queue=QUEUE_NAME,\n", + " ppn=PPN,\n", + ")\n", + "print(f\"Using cluster: {compute.cluster.hostname}, queue: {QUEUE_NAME}, ppn: {PPN}\")" + ] + }, + { + "cell_type": "markdown", + "id": "30", + "metadata": {}, + "source": [ + "## 6. Create the jobs with material and workflow configuration\n", + "### 6.1. Create one job per material" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "31", + "metadata": {}, + "outputs": [], + "source": [ + "from mat3ra.notebooks_utils.job import create_job\n", + "from mat3ra.notebooks_utils.ui import display_JSON\n", + "\n", + "jobs = []\n", + "for saved_material in saved_materials:\n", + " job_response = create_job(\n", + " api_client=client,\n", + " materials=[saved_material],\n", + " workflow=saved_workflow or workflow,\n", + " project_id=project_id,\n", + " owner_id=ACCOUNT_ID,\n", + " prefix=f\"{MY_WORKFLOW_NAME} {saved_material.formula} {timestamp}\",\n", + " compute=compute.to_dict(),\n", + " )\n", + " jobs.append(job_response if not isinstance(job_response, list) else job_response[0])\n", + "\n", + "job_ids = [job[\"_id\"] for job in jobs]\n", + "print(f\"✅ Created {len(job_ids)} jobs: {job_ids}\")\n", + "display_JSON(jobs[0])" + ] + }, + { + "cell_type": "markdown", + "id": "32", + "metadata": {}, + "source": [ + "## 7. Submit the jobs and monitor the status" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "33", + "metadata": {}, + "outputs": [], + "source": [ + "from mat3ra.notebooks_utils.api.job import submit_jobs\n", + "\n", + "submit_jobs(client.jobs, job_ids)\n", + "print(f\"✅ Submitted {len(job_ids)} jobs successfully!\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "34", + "metadata": {}, + "outputs": [], + "source": [ + "from mat3ra.notebooks_utils.api.job import wait_for_jobs_to_finish_async\n", + "\n", + "await wait_for_jobs_to_finish_async(client.jobs, job_ids, poll_interval=POLL_INTERVAL)" + ] + }, + { + "cell_type": "markdown", + "id": "35", + "metadata": {}, + "source": [ + "## 8. Retrieve results\n", + "\n", + "Each job's standard output is the script's own output. Lines that parse as JSON become table\n", + "columns; everything else is shown as printed." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "36", + "metadata": {}, + "outputs": [], + "source": [ + "import pandas as pd\n", + "\n", + "from mat3ra.notebooks_utils.io import read_from_url\n", + "\n", + "\n", + "# The runner names a unit's stdout after the unit, so this is the execution unit's own output -\n", + "# not whatever .out file the script itself may have written next to it.\n", + "STDOUT_FILENAME = \"custom_script.out\"\n", + "\n", + "\n", + "async def read_job_stdout(job_id):\n", + " \"\"\"Contents of the execution unit's stdout, or why the job produced none.\"\"\"\n", + " files = client.jobs.list_files(job_id)\n", + " stdout_file = next((file for file in files if file[\"key\"].rsplit(\"/\", 1)[-1] == STDOUT_FILENAME), None)\n", + " if stdout_file is None:\n", + " job = client.jobs.get(job_id)\n", + " errors = job.get(\"compute\", {}).get(\"errors\", [])\n", + " return f\"No output. Job status: {job['status']}. {json.dumps(errors, indent=2)}\"\n", + " return await read_from_url(stdout_file[\"signedUrl\"])\n", + "\n", + "\n", + "results = []\n", + "for saved_material, job in zip(saved_materials, jobs):\n", + " stdout = await read_job_stdout(job[\"_id\"])\n", + " print(f\"--- {job['name']} ---\\n{stdout}\")\n", + " for line in stdout.splitlines():\n", + " try:\n", + " results.append({\"material\": saved_material.name, **json.loads(line)})\n", + " except json.JSONDecodeError:\n", + " continue\n", + "\n", + "pd.DataFrame(results)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "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.11.2" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/other/materials_designer/workflows/custom_shell_calculation.ipynb b/other/materials_designer/workflows/custom_shell_calculation.ipynb new file mode 100644 index 000000000..b6377afc5 --- /dev/null +++ b/other/materials_designer/workflows/custom_shell_calculation.ipynb @@ -0,0 +1,705 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": {}, + "source": [ + "# Custom Shell Calculation\n", + "\n", + "Run your own shell script against one or more materials on the Mat3ra platform — the shell twin of\n", + "the Custom Python Calculation notebook. The script and any files it needs are uploaded to your\n", + "object storage folder, a workflow fetches them onto the compute node alongside the material, and\n", + "whatever the script prints comes back as the result. A shell script can invoke any application\n", + "installed on the node; the default example runs a Quantum ESPRESSO band structure for silicon with\n", + "an uploaded pseudopotential.\n", + "\n", + "

Usage

\n", + "\n", + "1. Set the materials and any data files the script reads in cell 1.2. below.\n", + "1. Replace the script in cell 1.3. with your own (or keep the default one).\n", + "1. Click \"Run\" > \"Run All\" to run all cells.\n", + "1. Wait for the jobs to complete.\n", + "1. Scroll down to view the results.\n", + "\n", + "## Summary\n", + "\n", + "1. Set up the environment and parameters: install packages (JupyterLite only) and configure the\n", + " materials, the script, its data files, compute resources and job.\n", + "1. Authenticate and initialize API client: authenticate via browser, initialize the client, then\n", + " select account and project.\n", + "1. Create materials: materials are read from the `../uploads` folder — place files there manually or\n", + " run a material creation notebook first. If a material is not found by name, Standata is used as a\n", + " fallback. Each material is then saved to the platform.\n", + "1. Create workflow: upload the script and its data files, then assemble a workflow that fetches\n", + " them, fetches the material, and runs the script. Optionally save the workflow to the collection.\n", + "1. Configure compute: get list of clusters and create compute configuration with selected cluster,\n", + " queue, and number of processors.\n", + "1. Create one job per material from the material, workflow, project and compute configuration.\n", + "1. Submit the jobs and monitor the status: submit and wait for completion.\n", + "1. Retrieve results: read each job's standard output and display the values it printed.\n", + "\n", + "## How the script receives its inputs\n", + "\n", + "Everything lands in the job's working directory, so the script reads it all by **relative path**:\n", + "\n", + "| File | Written by | Contents |\n", + "| --- | --- | --- |\n", + "| `material.json` | the workflow | the job's material, as stored on the platform |\n", + "| your data files | this notebook | uploaded verbatim from `USER_ASSET_FILES` |\n", + "\n", + "The script runs in a shell where `module` is available, so node-side applications load the same way\n", + "they do in a command-line job (e.g. `module add espresso`). The script's standard output is the\n", + "result. Print JSON and this notebook renders it as a table." + ] + }, + { + "cell_type": "markdown", + "id": "1", + "metadata": {}, + "source": [ + "## 1. Set up the environment and parameters\n", + "### 1.1. Install packages (JupyterLite)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "from mat3ra.notebooks_utils.packages import install_packages\n", + "\n", + "await install_packages(\"made|api_examples\")" + ] + }, + { + "cell_type": "markdown", + "id": "3", + "metadata": {}, + "source": [ + "### 1.2. Set parameters and configurations for the workflow and job\n", + "\n", + "`USER_ASSET_FILES` names data files the script opens, of any type. Put them in the `../uploads`\n", + "folder first — drag them into the JupyterLite file browser — and this notebook uploads them\n", + "alongside the script." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "from datetime import datetime\n", + "\n", + "from mat3ra.ide.compute import QueueName\n", + "\n", + "# 2. Auth and organization parameters\n", + "# Set organization name to use it as the owner, otherwise your personal account is used\n", + "ORGANIZATION_NAME = None\n", + "\n", + "# 3. Material parameters\n", + "FOLDER = \"../uploads\"\n", + "MATERIAL_NAMES = [\"Silicon\"] # One job is created per material\n", + "\n", + "# 4. Script parameters\n", + "USER_ASSET_FILES = [] # data files the script opens, taken from FOLDER, e.g. a pseudopotential\n", + "\n", + "# 5. Workflow parameters\n", + "WORKFLOW_SEARCH_TERM = \"custom_script.json\" # Search term for Workflows Standata\n", + "APPLICATION_NAME = \"shell\"\n", + "MY_WORKFLOW_NAME = \"Custom Shell Calculation\"\n", + "save_to_collection = True\n", + "\n", + "# 6. Compute parameters\n", + "CLUSTER_NAME = None # specify full or partial name i.e. \"cluster-001\" to select\n", + "QUEUE_NAME = QueueName.D\n", + "PPN = 1\n", + "\n", + "# 7. Job parameters\n", + "timestamp = datetime.now().strftime(\"%Y-%m-%d %H:%M\")\n", + "POLL_INTERVAL = 30 # seconds" + ] + }, + { + "cell_type": "markdown", + "id": "5", + "metadata": {}, + "source": [ + "### 1.3. Set the script to run\n", + "\n", + "This is the calculation. It runs on the compute node with `material.json` and your data files\n", + "beside it, and whatever it prints becomes the result. Replace it with your own.\n", + "\n", + "The default example runs a Quantum ESPRESSO band structure for silicon: it builds the pw.x inputs\n", + "from `material.json`, runs an SCF and a bands step, and prints the band edges as JSON. It reads the\n", + "pseudopotential from the platform's library; to use your own, put the file in `../uploads`, list it\n", + "in `USER_ASSET_FILES`, and set `PSEUDO_DIR` in the script to the working directory. The script is uploaded as a file and fetched onto the node, never inlined into\n", + "the workflow, so its contents reach the shell exactly as written." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "USER_SCRIPT = r\"\"\"#!/bin/bash\n", + "# Runs a Quantum ESPRESSO band structure for the job's material.\n", + "#\n", + "# The pseudopotential comes from the platform's library. To use your own instead, put it in\n", + "# uploads, list it in USER_ASSET_FILES, and point PSEUDO_DIR at the working directory (\"./\").\n", + "PSEUDO_DIR=\"/export/share/pseudo/si/gga/pbe/gbrv/1.0/us/\"\n", + "PSEUDO_FILES='{\"Si\": \"si_pbe_gbrv_1.0.upf\"}'\n", + "\n", + "module add espresso\n", + "\n", + "# Build the pw.x inputs from the job's material.\n", + "export PSEUDO_DIR PSEUDO_FILES\n", + "python3 - <<'BUILD_INPUTS_EOF'\n", + "import json\n", + "import math\n", + "\n", + "material = json.load(open(\"material.json\"))\n", + "lattice = material[\"lattice\"]\n", + "a, b, c = lattice[\"a\"], lattice[\"b\"], lattice[\"c\"]\n", + "alpha, beta, gamma = (math.radians(lattice[key]) for key in (\"alpha\", \"beta\", \"gamma\"))\n", + "c_x = c * math.cos(beta)\n", + "c_y = c * (math.cos(alpha) - math.cos(beta) * math.cos(gamma)) / math.sin(gamma)\n", + "vectors = [\n", + " [a, 0.0, 0.0],\n", + " [b * math.cos(gamma), b * math.sin(gamma), 0.0],\n", + " [c_x, c_y, math.sqrt(max(c**2 - c_x**2 - c_y**2, 0.0))],\n", + "]\n", + "elements = [element[\"value\"] for element in material[\"basis\"][\"elements\"]]\n", + "coordinates = [point[\"value\"] for point in material[\"basis\"][\"coordinates\"]]\n", + "\n", + "import os\n", + "MASSES = {\"Si\": 28.0855}\n", + "PSEUDOS = json.loads(os.environ[\"PSEUDO_FILES\"])\n", + "PSEUDO_DIR = os.environ[\"PSEUDO_DIR\"]\n", + "\n", + "cell = \"\\n\".join(\" \".join(f\"{value:.9f}\" for value in vector) for vector in vectors)\n", + "positions = \"\\n\".join(f\"{el} \" + \" \".join(f\"{x:.9f}\" for x in xyz) for el, xyz in zip(elements, coordinates))\n", + "species = \"\\n\".join(f\"{el} {MASSES[el]} {PSEUDOS[el]}\" for el in sorted(set(elements)))\n", + "\n", + "common = f'''&SYSTEM\n", + " ibrav = 0\n", + " nat = {len(elements)}\n", + " ntyp = {len(set(elements))}\n", + " ecutwfc = 40\n", + " ecutrho = 200\n", + " occupations = 'fixed'\n", + " nbnd = 8\n", + "/\n", + "&ELECTRONS\n", + "/\n", + "ATOMIC_SPECIES\n", + "{species}\n", + "CELL_PARAMETERS angstrom\n", + "{cell}\n", + "ATOMIC_POSITIONS crystal\n", + "{positions}\n", + "'''\n", + "\n", + "with open(\"pw_scf.in\", \"w\") as f:\n", + " f.write(f\"&CONTROL\\n calculation = 'scf'\\n pseudo_dir = '{PSEUDO_DIR}'\\n outdir = './outdir'\\n/\\n\")\n", + " f.write(common)\n", + " f.write(\"K_POINTS automatic\\n6 6 6 0 0 0\\n\")\n", + "\n", + "with open(\"pw_bands.in\", \"w\") as f:\n", + " f.write(f\"&CONTROL\\n calculation = 'bands'\\n pseudo_dir = '{PSEUDO_DIR}'\\n outdir = './outdir'\\n/\\n\")\n", + " f.write(common)\n", + " f.write(\"K_POINTS crystal_b\\n3\\n0.5 0.5 0.5 20\\n0.0 0.0 0.0 20\\n0.5 0.0 0.5 20\\n\")\n", + "BUILD_INPUTS_EOF\n", + "\n", + "# Run the calculation. EXEC_CMD is (conditionally) set by the module.\n", + "mpirun -np $PBS_NP $EXEC_CMD pw.x -in pw_scf.in > pw_scf.out\n", + "mpirun -np $PBS_NP $EXEC_CMD pw.x -in pw_bands.in > pw_bands.out\n", + "\n", + "# Proof and results: QE's own log names the pseudopotential file it read\n", + "# (the filename wraps onto the next line, hence -A 1).\n", + "grep -A 1 \"read from file\" pw_scf.out\n", + "grep \"highest occupied\" pw_scf.out\n", + "\n", + "# Band edges at Gamma from the bands run, as JSON for the results table.\n", + "python3 - <<'PARSE_BANDS_EOF'\n", + "import json\n", + "import re\n", + "\n", + "text = open(\"pw_bands.out\").read()\n", + "block = text[text.index(\"End of band structure calculation\"):]\n", + "gamma = re.search(r\"k = 0\\.0000 0\\.0000 0\\.0000[^\\n]*\\n\\n([\\s\\S]*?)\\n\\n\", block).group(1)\n", + "energies = sorted(float(value) for value in re.findall(r\"-?\\d+\\.\\d+\", gamma))\n", + "print(json.dumps({\n", + " \"gamma_homo_ev\": energies[3],\n", + " \"gamma_lumo_ev\": energies[4],\n", + " \"gamma_direct_gap_ev\": round(energies[4] - energies[3], 4),\n", + "}))\n", + "PARSE_BANDS_EOF\n", + "\"\"\"" + ] + }, + { + "cell_type": "markdown", + "id": "7", + "metadata": {}, + "source": [ + "## 2. Authenticate and initialize API client\n", + "### 2.1. Authenticate\n", + "Authenticate in the browser and have credentials stored in environment variable \"OIDC_ACCESS_TOKEN\"." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "metadata": {}, + "outputs": [], + "source": [ + "from mat3ra.notebooks_utils.auth import authenticate\n", + "\n", + "\n", + "await authenticate()" + ] + }, + { + "cell_type": "markdown", + "id": "9", + "metadata": {}, + "source": [ + "### 2.2. Initialize API Client" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10", + "metadata": {}, + "outputs": [], + "source": [ + "from mat3ra.api_client import APIClient\n", + "\n", + "client = APIClient.authenticate()\n", + "client" + ] + }, + { + "cell_type": "markdown", + "id": "11", + "metadata": {}, + "source": [ + "### 2.3. Select account to work under" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12", + "metadata": {}, + "outputs": [], + "source": [ + "client.list_accounts()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "13", + "metadata": {}, + "outputs": [], + "source": [ + "selected_account = client.my_account\n", + "\n", + "if ORGANIZATION_NAME:\n", + " selected_account = client.get_account(name=ORGANIZATION_NAME)\n", + "\n", + "ACCOUNT_ID = selected_account.id\n", + "print(f\"Using account: {selected_account.name} ({ACCOUNT_ID})\")" + ] + }, + { + "cell_type": "markdown", + "id": "14", + "metadata": {}, + "source": [ + "### 2.4. Select project" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "15", + "metadata": {}, + "outputs": [], + "source": [ + "projects = client.projects.list({\"isDefault\": True, \"owner._id\": ACCOUNT_ID})\n", + "project_id = projects[0][\"_id\"]\n", + "print(f\"✅ Using project: {projects[0]['name']} ({project_id})\")" + ] + }, + { + "cell_type": "markdown", + "id": "16", + "metadata": {}, + "source": [ + "## 3. Create materials\n", + "### 3.1. Load materials from local files (or Standata)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "17", + "metadata": {}, + "outputs": [], + "source": [ + "from mat3ra.made.material import Material\n", + "from mat3ra.standata.materials import Materials\n", + "from mat3ra.notebooks_utils.ipython.entity.material.visualize import visualize_materials as visualize\n", + "from mat3ra.notebooks_utils.material import load_material_from_folder\n", + "\n", + "materials = [\n", + " load_material_from_folder(FOLDER, name) or Material.create(Materials.get_by_name_first_match(name))\n", + " for name in MATERIAL_NAMES\n", + "]\n", + "visualize([{\"material\": material, \"title\": material.name} for material in materials])" + ] + }, + { + "cell_type": "markdown", + "id": "18", + "metadata": {}, + "source": [ + "### 3.2. Save materials to the platform" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "19", + "metadata": {}, + "outputs": [], + "source": [ + "from mat3ra.notebooks_utils.core.entity.material.api import get_or_create_material\n", + "\n", + "saved_materials = [\n", + " Material.create(get_or_create_material(client, material, ACCOUNT_ID)) for material in materials\n", + "]" + ] + }, + { + "cell_type": "markdown", + "id": "20", + "metadata": {}, + "source": [ + "## 4. Create workflow and set its parameters\n", + "### 4.1. Upload the script and its data files\n", + "\n", + "The files go to your account's object storage folder (\"Dropbox\"), which the compute node reads them\n", + "from. Text travels inside the request; anything else — a model checkpoint, an archive, a large file\n", + "— is sent to a URL the platform signs for it, straight to storage, so there is no size or format\n", + "limit beyond what the browser can hold in memory." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "21", + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "\n", + "from mat3ra.notebooks_utils.core.entity.file.api import upload_files\n", + "\n", + "files_to_upload = {\"user_script.sh\": USER_SCRIPT}\n", + "for name in USER_ASSET_FILES:\n", + " if name in files_to_upload:\n", + " raise ValueError(f\"Rename '{name}' in USER_ASSET_FILES: this notebook already uploads a file by that name.\")\n", + " path = os.path.join(FOLDER, name)\n", + " if not os.path.exists(path):\n", + " raise FileNotFoundError(f\"'{name}' is listed in USER_ASSET_FILES but is not in {FOLDER}.\")\n", + " with open(path, \"rb\") as file:\n", + " content = file.read()\n", + " try:\n", + " files_to_upload[name] = content.decode(\"utf-8\") # text travels in the request body\n", + " except UnicodeDecodeError:\n", + " files_to_upload[name] = content # anything else goes to storage through a signed URL\n", + "\n", + "uploaded_files = upload_files(client, files_to_upload, ACCOUNT_ID)" + ] + }, + { + "cell_type": "markdown", + "id": "22", + "metadata": {}, + "source": [ + "### 4.2. Create workflow from standard workflows and preview it\n", + "\n", + "The `Custom Shell Script` workflow already carries the unit chain this needs: fetch the uploaded\n", + "files, fetch the material, run the script. Two things are filled in per job — the objects to fetch\n", + "and the runner that hands your script its inputs." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "23", + "metadata": {}, + "outputs": [], + "source": [ + "from mat3ra.standata.applications import ApplicationStandata\n", + "from mat3ra.ade.application import Application\n", + "from mat3ra.standata.workflows import WorkflowStandata\n", + "from mat3ra.wode.workflows import Workflow\n", + "from mat3ra.notebooks_utils.core.entity.file.api import to_object_storage_input\n", + "from mat3ra.notebooks_utils.core.entity.workflow.api import CUSTOM_SCRIPT_RUNNER_SH, set_execution_unit_input\n", + "from mat3ra.notebooks_utils.ipython.entity.workflow.visualize import visualize_workflow\n", + "\n", + "app_config = ApplicationStandata.get_by_name_first_match(APPLICATION_NAME)\n", + "app = Application(**app_config)\n", + "\n", + "workflow_config = WorkflowStandata.filter_by_application(app.name).get_by_name_first_match(WORKFLOW_SEARCH_TERM)\n", + "workflow_config[\"name\"] = MY_WORKFLOW_NAME\n", + "subworkflow = workflow_config[\"subworkflows\"][0]\n", + "subworkflow[\"name\"] = MY_WORKFLOW_NAME\n", + "units = {unit[\"name\"]: unit for unit in subworkflow[\"units\"]}\n", + "\n", + "units[\"io-user-files\"][\"input\"] = [to_object_storage_input(file) for file in uploaded_files]\n", + "\n", + "set_execution_unit_input(units[\"custom_script\"], \"hello_world.sh\", CUSTOM_SCRIPT_RUNNER_SH)\n", + "\n", + "workflow = Workflow.create(workflow_config)\n", + "visualize_workflow(workflow)" + ] + }, + { + "cell_type": "markdown", + "id": "24", + "metadata": {}, + "source": [ + "### 4.3. Save workflow to collection\n", + "\n", + "Saving it makes the workflow reusable: it stays in your collection pointing at the uploaded files,\n", + "so it can be run again — from the UI or another notebook — against any material. Re-running this\n", + "notebook overwrites the uploaded files in place, so a saved workflow picks up an edited script\n", + "without being rebuilt." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "25", + "metadata": {}, + "outputs": [], + "source": [ + "saved_workflow = None\n", + "if save_to_collection:\n", + " saved_workflow = Workflow.create(\n", + " client.workflows.create(workflow.to_dict_without_special_keys(), owner_id=ACCOUNT_ID)\n", + " )\n", + " print(f\"✅ Workflow saved to collection: {saved_workflow.id}\")" + ] + }, + { + "cell_type": "markdown", + "id": "26", + "metadata": {}, + "source": [ + "## 5. Create the compute configuration\n", + "### 5.1. Get list of clusters" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "27", + "metadata": {}, + "outputs": [], + "source": [ + "clusters = client.clusters.list()\n", + "print(f\"Available clusters: {[c['hostname'] for c in clusters]}\")" + ] + }, + { + "cell_type": "markdown", + "id": "28", + "metadata": {}, + "source": [ + "### 5.2. Create compute configuration for the jobs" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "29", + "metadata": {}, + "outputs": [], + "source": [ + "from mat3ra.ide.compute import Compute\n", + "\n", + "# Select cluster: use specified name if provided, otherwise use first available\n", + "if CLUSTER_NAME:\n", + " cluster = next((c for c in clusters if CLUSTER_NAME in c[\"hostname\"]), None)\n", + "else:\n", + " cluster = clusters[0]\n", + "\n", + "compute = Compute(\n", + " cluster=cluster,\n", + " queue=QUEUE_NAME,\n", + " ppn=PPN,\n", + ")\n", + "print(f\"Using cluster: {compute.cluster.hostname}, queue: {QUEUE_NAME}, ppn: {PPN}\")" + ] + }, + { + "cell_type": "markdown", + "id": "30", + "metadata": {}, + "source": [ + "## 6. Create the jobs with material and workflow configuration\n", + "### 6.1. Create one job per material" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "31", + "metadata": {}, + "outputs": [], + "source": [ + "from mat3ra.notebooks_utils.job import create_job\n", + "from mat3ra.notebooks_utils.ui import display_JSON\n", + "\n", + "jobs = []\n", + "for saved_material in saved_materials:\n", + " job_response = create_job(\n", + " api_client=client,\n", + " materials=[saved_material],\n", + " workflow=saved_workflow or workflow,\n", + " project_id=project_id,\n", + " owner_id=ACCOUNT_ID,\n", + " prefix=f\"{MY_WORKFLOW_NAME} {saved_material.formula} {timestamp}\",\n", + " compute=compute.to_dict(),\n", + " )\n", + " jobs.append(job_response if not isinstance(job_response, list) else job_response[0])\n", + "\n", + "job_ids = [job[\"_id\"] for job in jobs]\n", + "print(f\"✅ Created {len(job_ids)} jobs: {job_ids}\")\n", + "display_JSON(jobs[0])" + ] + }, + { + "cell_type": "markdown", + "id": "32", + "metadata": {}, + "source": [ + "## 7. Submit the jobs and monitor the status" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "33", + "metadata": {}, + "outputs": [], + "source": [ + "from mat3ra.notebooks_utils.api.job import submit_jobs\n", + "\n", + "submit_jobs(client.jobs, job_ids)\n", + "print(f\"✅ Submitted {len(job_ids)} jobs successfully!\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "34", + "metadata": {}, + "outputs": [], + "source": [ + "from mat3ra.notebooks_utils.api.job import wait_for_jobs_to_finish_async\n", + "\n", + "await wait_for_jobs_to_finish_async(client.jobs, job_ids, poll_interval=POLL_INTERVAL)" + ] + }, + { + "cell_type": "markdown", + "id": "35", + "metadata": {}, + "source": [ + "## 8. Retrieve results\n", + "\n", + "Each job's standard output is the script's own output. Lines that parse as JSON become table\n", + "columns; everything else is shown as printed." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "36", + "metadata": {}, + "outputs": [], + "source": [ + "import pandas as pd\n", + "\n", + "from mat3ra.notebooks_utils.io import read_from_url\n", + "\n", + "\n", + "# The runner names a unit's stdout after the unit, so this is the execution unit's own output -\n", + "# not whatever .out file the script itself may have written next to it.\n", + "STDOUT_FILENAME = \"custom_script.out\"\n", + "\n", + "\n", + "async def read_job_stdout(job_id):\n", + " \"\"\"Contents of the execution unit's stdout, or why the job produced none.\"\"\"\n", + " files = client.jobs.list_files(job_id)\n", + " stdout_file = next((file for file in files if file[\"key\"].rsplit(\"/\", 1)[-1] == STDOUT_FILENAME), None)\n", + " if stdout_file is None:\n", + " job = client.jobs.get(job_id)\n", + " errors = job.get(\"compute\", {}).get(\"errors\", [])\n", + " return f\"No output. Job status: {job['status']}. {json.dumps(errors, indent=2)}\"\n", + " return await read_from_url(stdout_file[\"signedUrl\"])\n", + "\n", + "\n", + "results = []\n", + "for saved_material, job in zip(saved_materials, jobs):\n", + " stdout = await read_job_stdout(job[\"_id\"])\n", + " print(f\"--- {job['name']} ---\\n{stdout}\")\n", + " for line in stdout.splitlines():\n", + " try:\n", + " results.append({\"material\": saved_material.name, **json.loads(line)})\n", + " except json.JSONDecodeError:\n", + " continue\n", + "\n", + "pd.DataFrame(results)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "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.11.2" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/pyproject.toml b/pyproject.toml index bd6adf9b2..345457053 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -71,6 +71,9 @@ tests = [ "pytest", "pytest-asyncio", "pytest-cov", + # notebooks_utils.io imports IPython eagerly (via pyodide.io), so anything reaching it for + # get_data/set_data - core.entity.material.io, for one - needs IPython importable under test. + "ipython>=8.0", "mat3ra-notebooks-utils[workflows]", ] docs = [ diff --git a/src/py/mat3ra/notebooks_utils/core/entity/file/__init__.py b/src/py/mat3ra/notebooks_utils/core/entity/file/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/py/mat3ra/notebooks_utils/core/entity/file/api.py b/src/py/mat3ra/notebooks_utils/core/entity/file/api.py new file mode 100644 index 000000000..c4847f8bd --- /dev/null +++ b/src/py/mat3ra/notebooks_utils/core/entity/file/api.py @@ -0,0 +1,133 @@ +import json +from typing import Dict, List, Union + +from mat3ra.api_client import APIClient +from mat3ra.api_client.endpoints import BaseEndpoint +from mat3ra.notebooks_utils.primitive.environment import is_pyodide_environment + +# All of these land in the job's working directory after the IO unit has fetched the uploaded +# files, so an upload under any of them is overwritten before the user's script runs: the execution +# unit renders `script.py` and `requirements.txt`, and the runner it renders writes `material.json`. +# Keep in step with CUSTOM_SCRIPT_RUNNER / CUSTOM_SCRIPT_RUNNER_SH in ../workflow/api.py. +# `hello_world.sh` is the on-disk name of the shell runner (the shell flavor's input name). +RESERVED_FILENAMES = ("script.py", "requirements.txt", "material.json", "hello_world.sh") + + +def _files_endpoint(api_client: APIClient) -> BaseEndpoint: + """ + A raw endpoint for `/files`, which mat3ra-api-client does not model yet. + + Args: + api_client (APIClient): API client instance carrying the authorization context. + + Returns: + BaseEndpoint: Endpoint bound to the same host, version and credentials as the client. + """ + return BaseEndpoint( + api_client.host, + api_client.port, + version=api_client.version, + secure=api_client.secure, + auth=api_client.auth, + ) + + +def upload_files(api_client: APIClient, files: Dict[str, Union[str, bytes]], account_id: str) -> List[dict]: + """ + Uploads files to the account's object storage ("Dropbox") folder. + + Text travels inside the request body (`POST /files`). Bytes - a model checkpoint, an archive, + anything that is not UTF-8 text or is large - go through a URL the platform signs for the + purpose (`POST /files/signed-urls`) and are PUT to storage directly, so neither the JSON body + limit nor the text encoding applies. + + Args: + api_client (APIClient): API client instance carrying the authorization context. + files (dict): File name (relative to the account folder) mapped to its content, `str` for + text or `bytes` for anything else. + account_id (str): Account to upload under. + + Returns: + list[dict]: One cloud file record per upload, with `key`, `size`, `bucket`, `region` + and `provider`. + + Raises: + ValueError: If a file name collides with one the execution unit writes itself. + """ + reserved = [name for name in files if name.split("/")[-1] in RESERVED_FILENAMES] + if reserved: + raise ValueError(f"Rename {reserved}: {RESERVED_FILENAMES} are written by the workflow itself.") + + endpoint = _files_endpoint(api_client) + headers = endpoint.get_headers(api_client.auth.account_id or "", api_client.auth.auth_token or "") + + uploaded = [] + for name, content in files.items(): + if isinstance(content, bytes): + record = _upload_bytes(endpoint, headers, name, content, account_id) + else: + payload = {"name": name, "body": content, "accountId": account_id} + record = endpoint.request("POST", "files", data=json.dumps(payload), headers=headers) + print(f"⬆️ Uploaded {record['key']} ({record['size']} bytes)") + uploaded.append(record) + return uploaded + + +def _upload_bytes(endpoint: BaseEndpoint, headers: dict, name: str, content: bytes, account_id: str) -> dict: + """Asks the platform for a signed PUT URL for `name`, PUTs the bytes, returns a file record.""" + payload = {"names": [name], "operation": "putObject", "accountId": account_id} + [target] = endpoint.request("POST", "files/signed-urls", data=json.dumps(payload), headers=headers) + _put(target["signedUrl"], content) + return { + "name": name, + "key": target["key"], + "size": len(content), + "bucket": target["bucket"], + "region": target["region"], + "provider": target["provider"], + } + + +def _put(url: str, data: bytes) -> None: + """PUTs bytes to a signed URL - from the browser in JupyterLite, from the process elsewhere.""" + if not is_pyodide_environment(): + from urllib.request import Request, urlopen + + put = Request(url, data=data, method="PUT") + put.add_header("Content-Type", "application/octet-stream") + with urlopen(put) as response: + response.read() + return + + from js import XMLHttpRequest # type: ignore[import-not-found] + from pyodide.ffi import to_js # type: ignore[import-not-found] + + request = XMLHttpRequest.new() + request.open("PUT", url, False) + request.send(to_js(data)) + if request.status < 200 or request.status >= 300: + raise RuntimeError(f"PUT to storage failed with HTTP {request.status}: {request.responseText}") + + +def to_object_storage_input(cloud_file: dict) -> dict: + """ + Converts an upload record into an `object_storage` input item for a workflow IO unit. + + Args: + cloud_file (dict): A record returned by `upload_files`. + + Returns: + dict: IO unit input item. The runner fetches it into the job's working directory under + `basename`, which is what makes the user script's relative paths resolve. + """ + return { + "type": "object_storage", + "basename": cloud_file["key"].split("/")[-1], + "pathname": "", + "objectData": { + "NAME": cloud_file["key"], + "PROVIDER": cloud_file["provider"], + "CONTAINER": cloud_file["bucket"], + "REGION": cloud_file["region"], + }, + } diff --git a/src/py/mat3ra/notebooks_utils/core/entity/material/io.py b/src/py/mat3ra/notebooks_utils/core/entity/material/io.py index 5c60dd5fd..11866eee9 100644 --- a/src/py/mat3ra/notebooks_utils/core/entity/material/io.py +++ b/src/py/mat3ra/notebooks_utils/core/entity/material/io.py @@ -1,7 +1,7 @@ import inspect import json import os -from typing import Any, List, Optional +from typing import Any, Dict, List, Optional from mat3ra.made.material import Material from mat3ra.made.tools.build_components import MaterialWithBuildMetadata @@ -64,6 +64,28 @@ def set_materials(materials: List[Any], folder_path: str = UPLOADS_FOLDER): ) +def _create_material_or_none(config: Dict[str, Any]) -> Optional[Any]: + """ + Builds a material from a config, or returns None when the config is not one. + + The folder holds whatever the user put there — a data file a script reads is as likely as a + material — so a config no material class accepts is skipped instead of failing the whole load. + Build metadata is preferred per file, so one plain material does not strip it from the rest. + + Args: + config (dict): Parsed contents of a JSON file from the folder. + + Returns: + Material | MaterialWithBuildMetadata | None: The material, or None if the config is not one. + """ + for material_cls in (MaterialWithBuildMetadata, Material): + try: + return material_cls.create(config) + except Exception: + continue + return None + + def load_materials_from_folder(folder_path: Optional[str] = None, verbose: bool = True) -> List[Any]: """ Load materials from the specified folder or from the UPLOADS_FOLDER by default. @@ -84,7 +106,6 @@ def load_materials_from_folder(folder_path: Optional[str] = None, verbose: bool data_from_host = [] try: - index = 0 for filename in sorted(os.listdir(folder_path)): if filename.endswith(".json"): file_path = os.path.join(folder_path, filename) @@ -98,18 +119,23 @@ def load_materials_from_folder(folder_path: Optional[str] = None, verbose: bool force_verbose=verbose, ) continue - name = os.path.splitext(filename)[0] - log(f"{index}: {name}", SeverityLevelEnum.INFO, force_verbose=verbose) - index += 1 - data_from_host.append(data) + data_from_host.append((os.path.splitext(filename)[0], data)) except FileNotFoundError: log(f"No data found in the '{folder_path}' folder.", SeverityLevelEnum.ERROR, force_verbose=verbose) return [] - try: - materials = [MaterialWithBuildMetadata.create(item) for item in data_from_host] - except Exception: - materials = [Material.create(item) for item in data_from_host] + materials: List[Any] = [] + for name, item in data_from_host: + material = _create_material_or_none(item) + if material is None: + log( + f"Skipping '{name}.json': not a material.", + SeverityLevelEnum.WARNING, + force_verbose=verbose, + ) + continue + log(f"{len(materials)}: {name}", SeverityLevelEnum.INFO, force_verbose=verbose) + materials.append(material) if materials: log( diff --git a/src/py/mat3ra/notebooks_utils/core/entity/workflow/api.py b/src/py/mat3ra/notebooks_utils/core/entity/workflow/api.py index 8694cfced..fc932d520 100644 --- a/src/py/mat3ra/notebooks_utils/core/entity/workflow/api.py +++ b/src/py/mat3ra/notebooks_utils/core/entity/workflow/api.py @@ -37,3 +37,56 @@ def copy_bank_workflow_by_system_name(endpoint: BankWorkflowEndpoints, system_na """ bank_workflow_id = endpoint.list({"systemName": system_name})[0]["_id"] return endpoint.copy(bank_workflow_id, account_id)["_id"] + + +# Written by the runner into the job's working directory, next to the files the IO unit fetched. +# The filename it writes is also listed in RESERVED_FILENAMES in ../file/api.py, which refuses an +# upload that would be overwritten by it; change both together. +CUSTOM_SCRIPT_RUNNER = '''import json + +with open("material.json", "w") as file: + json.dump(json.loads(r"""{{ MATERIAL | default({}) | tojson }}"""), file) + +with open("user_script.py") as file: + exec(compile(file.read(), "user_script.py", "exec")) +''' + + +# The shell twin of CUSTOM_SCRIPT_RUNNER: rendered as `hello_world.sh` (the shell flavor's input +# name), it writes the job's material next to the uploaded files and hands control to the user's +# script. The script is *sourced*, not run in a subshell, so `module` stays available - rupy runs +# the runner the same way. Also listed in RESERVED_FILENAMES in ../file/api.py. +CUSTOM_SCRIPT_RUNNER_SH = """#!/bin/bash +cat > material.json <<'MATERIAL_JSON_EOF' +{{ MATERIAL | default({}) | tojson }} +MATERIAL_JSON_EOF + +. ./user_script.sh +""" + + +def set_execution_unit_input(unit: dict, template_name: str, content: str) -> None: + """ + Replaces one input file of an execution unit with fixed content. + + `isManuallyChanged` stops the platform re-rendering the file at job creation, which is what + lets the runner keep its `{{ ... }}` placeholders for the compute node to resolve, and keeps a + user's requirements verbatim. Inputs are matched by template name rather than position, so + adding a file to the flavor cannot silently write content into the wrong one. + + Args: + unit (dict): Execution unit config from a workflow. + template_name (str): Name of the input file to replace, e.g. "script.py". + content (str): Content to write. + + Raises: + KeyError: If the unit has no input with that template name. + """ + for unit_input in unit["input"]: + if unit_input["template"]["name"] == template_name: + unit_input["template"]["content"] = content + unit_input["rendered"] = content + unit_input["isManuallyChanged"] = True + return + available = [unit_input["template"]["name"] for unit_input in unit["input"]] + raise KeyError(f"No input named '{template_name}' in unit '{unit.get('name')}'. Available: {available}") diff --git a/tests/py/unit/core/entity/test_file_api.py b/tests/py/unit/core/entity/test_file_api.py new file mode 100644 index 000000000..7637aeec3 --- /dev/null +++ b/tests/py/unit/core/entity/test_file_api.py @@ -0,0 +1,154 @@ +import json +from typing import Any, Dict +from unittest.mock import MagicMock + +import pytest +from mat3ra.notebooks_utils.core.entity.file.api import RESERVED_FILENAMES, to_object_storage_input, upload_files + +ACCOUNT_ID = "account-1" + +CLOUD_FILE: Dict[str, Any] = { + "key": "my-account/user_script.py", + "size": 14, + "bucket": "test-bucket", + "region": "us-west-2", + "provider": "aws", + "lastModified": 1609891535, + "name": "user_script.py", +} + + +def _client(*records): + """An APIClient whose files endpoint returns the given records, one per request.""" + client = MagicMock() + client.host, client.port, client.version, client.secure = "localhost", 3000, "2018-10-01", False + client.auth.account_id, client.auth.auth_token = ACCOUNT_ID, "token" + client.auth.access_token = None + return client + + +def test_upload_posts_name_and_body_per_file(monkeypatch): + requests = [] + + def request(method, path, data=None, headers=None): + requests.append((method, path, json.loads(data))) + return CLOUD_FILE + + endpoint = MagicMock() + endpoint.request.side_effect = request + endpoint.get_headers.return_value = {} + monkeypatch.setattr("mat3ra.notebooks_utils.core.entity.file.api._files_endpoint", lambda _: endpoint) + + uploaded = upload_files(_client(), {"user_script.py": "print(1)", "radii.json": "{}"}, ACCOUNT_ID) + + assert [r[:2] for r in requests] == [("POST", "files"), ("POST", "files")] + assert [r[2] for r in requests] == [ + {"name": "user_script.py", "body": "print(1)", "accountId": ACCOUNT_ID}, + {"name": "radii.json", "body": "{}", "accountId": ACCOUNT_ID}, + ] + assert uploaded == [CLOUD_FILE, CLOUD_FILE] + + +@pytest.mark.parametrize("name", RESERVED_FILENAMES) +def test_upload_refuses_a_name_the_job_would_overwrite(name, monkeypatch): + endpoint = MagicMock() + monkeypatch.setattr("mat3ra.notebooks_utils.core.entity.file.api._files_endpoint", lambda _: endpoint) + + with pytest.raises(ValueError, match=name): + upload_files(_client(), {name: "x"}, ACCOUNT_ID) + + endpoint.request.assert_not_called() + + +def test_object_storage_input_carries_every_field_the_runner_needs(): + """rupy requires all of NAME/PROVIDER/CONTAINER/REGION plus a basename, with no defaults.""" + assert to_object_storage_input(CLOUD_FILE) == { + "type": "object_storage", + "basename": "user_script.py", + "pathname": "", + "objectData": { + "NAME": "my-account/user_script.py", + "PROVIDER": "aws", + "CONTAINER": "test-bucket", + "REGION": "us-west-2", + }, + } + + +def test_object_storage_basename_is_the_name_the_script_opens(): + nested = {**CLOUD_FILE, "key": "my-account/assets/radii.json"} + + assert to_object_storage_input(nested)["basename"] == "radii.json" + + +def test_upload_refuses_the_shell_runner_name(): + """hello_world.sh is what the shell workflow's runner lands as - an upload would be clobbered.""" + with pytest.raises(ValueError, match="hello_world.sh"): + upload_files(None, {"hello_world.sh": "echo hi"}, "account") + + +def test_bytes_go_through_a_signed_put_url(monkeypatch): + """Bytes never enter a JSON body: the helper asks for a putObject URL and PUTs the content.""" + from mat3ra.notebooks_utils.core.entity.file import api as file_api + + calls: list = [] + puts: list = [] + endpoint = MagicMock() + endpoint.get_headers.return_value = {} + + def request(method, path, data=None, headers=None): + calls.append((method, path, json.loads(data))) + return [ + { + "key": "user-abc/model.pt", + "signedUrl": "https://s3/put?sig", + "bucket": "b", + "region": "r", + "provider": "aws", + } + ] + + endpoint.request.side_effect = request + monkeypatch.setattr(file_api, "_files_endpoint", lambda client: endpoint) + monkeypatch.setattr(file_api, "_put", lambda url, data: puts.append((url, data))) + + client = MagicMock() + client.auth.account_id = "acc" + client.auth.auth_token = "tok" + [record] = upload_files(client, {"model.pt": b"\x00\x01binary"}, "acc") + + assert calls == [ + ("POST", "files/signed-urls", {"names": ["model.pt"], "operation": "putObject", "accountId": "acc"}) + ] + assert puts == [("https://s3/put?sig", b"\x00\x01binary")] + assert record == { + "name": "model.pt", + "key": "user-abc/model.pt", + "size": 8, + "bucket": "b", + "region": "r", + "provider": "aws", + } + assert to_object_storage_input(record)["objectData"]["NAME"] == "user-abc/model.pt" + + +def test_text_still_travels_in_the_body(monkeypatch): + from mat3ra.notebooks_utils.core.entity.file import api as file_api + + endpoint = MagicMock() + endpoint.get_headers.return_value = {} + endpoint.request.return_value = { + "key": "user-abc/a.txt", + "size": 2, + "bucket": "b", + "region": "r", + "provider": "aws", + "name": "a.txt", + } + monkeypatch.setattr(file_api, "_files_endpoint", lambda client: endpoint) + client = MagicMock() + client.auth.account_id = "acc" + client.auth.auth_token = "tok" + upload_files(client, {"a.txt": "hi"}, "acc") + method, path = endpoint.request.call_args[0][:2] + assert (method, path) == ("POST", "files") diff --git a/tests/py/unit/core/entity/test_material_io.py b/tests/py/unit/core/entity/test_material_io.py new file mode 100644 index 000000000..fec281162 --- /dev/null +++ b/tests/py/unit/core/entity/test_material_io.py @@ -0,0 +1,55 @@ +import json +from typing import Any, Dict + +from mat3ra.notebooks_utils.core.entity.material.io import load_material_from_folder, load_materials_from_folder + +SILICON: Dict[str, Any] = { + "name": "Silicon", + "lattice": { + "a": 3.867, + "b": 3.867, + "c": 3.867, + "alpha": 60.0, + "beta": 60.0, + "gamma": 60.0, + "units": {"length": "angstrom", "angle": "degree"}, + "type": "FCC", + }, + "basis": { + "elements": [{"id": 0, "value": "Si"}, {"id": 1, "value": "Si"}], + "coordinates": [{"id": 0, "value": [0.0, 0.0, 0.0]}, {"id": 1, "value": [0.25, 0.25, 0.25]}], + "units": "crystal", + }, +} + +# The kind of file a user script reads - valid JSON, not a material. +RADII: Dict[str, Any] = {"Si": 1.11, "C": 0.76} + + +def _write(folder, name, payload): + path = folder / name + path.write_text(json.dumps(payload)) + return path + + +def test_non_material_json_is_skipped_not_fatal(tmp_path): + """A data file beside the materials must not fail the whole load.""" + _write(tmp_path, "silicon.json", SILICON) + _write(tmp_path, "radii.json", RADII) + + materials = load_materials_from_folder(str(tmp_path), verbose=False) + + assert [material.name for material in materials] == ["Silicon"] + + +def test_lookup_by_name_still_works_alongside_a_data_file(tmp_path): + _write(tmp_path, "silicon.json", SILICON) + _write(tmp_path, "radii.json", RADII) + + assert load_material_from_folder(str(tmp_path), "Silicon", verbose=False).name == "Silicon" + + +def test_folder_of_only_data_files_yields_nothing(tmp_path): + _write(tmp_path, "radii.json", RADII) + + assert load_materials_from_folder(str(tmp_path), verbose=False) == [] diff --git a/tests/py/unit/core/entity/test_workflow_api.py b/tests/py/unit/core/entity/test_workflow_api.py new file mode 100644 index 000000000..337db1851 --- /dev/null +++ b/tests/py/unit/core/entity/test_workflow_api.py @@ -0,0 +1,40 @@ +import pytest +from mat3ra.notebooks_utils.core.entity.workflow.api import set_execution_unit_input + + +def _unit(): + return { + "name": "custom_script", + "input": [ + {"template": {"name": "script.py", "content": "original"}}, + {"template": {"name": "requirements.txt", "content": ""}}, + ], + } + + +def test_replaces_the_named_input_and_marks_it_manually_changed(): + """isManuallyChanged is what stops the platform re-rendering the runner's placeholders.""" + unit = _unit() + set_execution_unit_input(unit, "requirements.txt", "numpy<2\n") + + changed = unit["input"][1] + assert changed["template"]["content"] == "numpy<2\n" + assert changed["rendered"] == "numpy<2\n" + assert changed["isManuallyChanged"] is True + assert unit["input"][0]["template"]["content"] == "original", "other inputs are untouched" + + +def test_matches_by_name_not_position(): + """A flavor may list its inputs in any order; position would write into the wrong file.""" + unit = _unit() + unit["input"].reverse() + set_execution_unit_input(unit, "script.py", "print(1)") + + by_name = {item["template"]["name"]: item for item in unit["input"]} + assert by_name["script.py"]["rendered"] == "print(1)" + assert "rendered" not in by_name["requirements.txt"] + + +def test_unknown_input_name_is_an_error_naming_what_is_available(): + with pytest.raises(KeyError, match="requirements.txt"): + set_execution_unit_input(_unit(), "nope.txt", "x")