Refactored NerscClient class with:

- sfapi_client class inheritance
- refactored ptycho specific code
- Modified 7012 config class
- sample jupyter notebook to submit a sample job with new NerscClient class

examples/sfapi_NerscClient_example.ipynb

@@ -0,0 +1,153 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from orchestration.nersc import NerscClient"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "client_id_path = \"../clientid.txt\"\n",
+    "sfapi_key_path = \"../sfapi_training.json\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "client = NerscClient(client_id_path, sfapi_key_path)\n",
+    "\n",
+    "user = client.user()\n",
+    "user"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "home_path = f\"/global/homes/{user.name[0]}/{user.name}\"\n",
+    "scratch_path = f\"/pscratch/sd/{user.name[0]}/{user.name}\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### The job script below will run a simple python program to generate random numbers\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "N = 15\n",
+    "\n",
+    "job_script = f\"\"\"#!/bin/bash\n",
+    "\n",
+    "#SBATCH -q debug\n",
+    "#SBATCH -A als\n",
+    "#SBATCH -N 1\n",
+    "#SBATCH -C cpu\n",
+    "#SBATCH -t 00:10:00\n",
+    "#SBATCH -J sfapi-demo\n",
+    "#SBATCH --exclusive\n",
+    "#SBATCH --output={scratch_path}/nerscClient-test/sfapi-demo-%j.out\n",
+    "#SBATCH --error={scratch_path}/nerscClient-test/sfapi-demo-%j.error\n",
+    "\n",
+    "module load python\n",
+    "# Prints N random numbers to form a normal disrobution\n",
+    "python -c \"import numpy as np; numbers = np.random.normal(size={N}); [print(n) for n in numbers]\"\n",
+    "\"\"\" "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(job_script)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "### Make sure our output folder is there for our data to go to\n",
+    "client.perlmutter.run(f\"mkdir -p {scratch_path}/nerscClient-test\")\n",
+    "# We can run ls on the directory to see that it was created\n",
+    "[output_dir] = client.perlmutter.ls(f\"{scratch_path}/nerscClient-demo\", directory=True)\n",
+    "\n",
+    "\n",
+    "# Check that the directory is there\n",
+    "output_dir.is_dir()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "### Submit the job and wait for the job to complete\n",
+    "job = client.perlmutter.submit_job(job_script)\n",
+    "print(job)\n",
+    "# Let's save the job id to use later \n",
+    "job_id = job.jobid\n",
+    "\n",
+    "print(f\"Waiting for job {job_id} to finish!\")\n",
+    "# Wait for the job to finish\n",
+    "job.complete()\n",
+    "print(\"Done!\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "job.state\n",
+    "print(job.state)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "env",
+   "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.9.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

orchestration/flows/bl7012/config.py

@@ -1,7 +1,7 @@
 from globus_sdk import TransferClient
 
 from orchestration import globus
-from orchestration.nersc import NerscClient
+from ptycho_nersc import NerscPtychoClient
 
 
 class Config7012:
@@ -14,7 +14,7 @@ def __init__(
         self.endpoints = globus.build_endpoints(config)
         self.apps = globus.build_apps(config)
         self.tc: TransferClient = globus.init_transfer_client(self.apps["als_transfer"])
-        self.nersc = NerscClient(
+        self.nersc = NerscPtychoClient(
             path_client_id,
             path_priv_key,
         )

orchestration/flows/bl7012/ptycho_nersc.py

@@ -0,0 +1,48 @@
+import json
+import logging
+import time
+
+from authlib.integrations.requests_client import OAuth2Session
+from authlib.oauth2.rfc7523 import PrivateKeyJWT
+
+from orchestration.ptycho_jobscript import (
+    get_job_script,
+    cdtool_args_string,
+    ptychocam_args_string,
+    cdtools_parms,
+    ptychocam_parms,
+)
+
+from orchestration.nersc import NerscClient
+
+class NerscPtychoClient(NerscClient):
+    def __init__(
+        self,
+        path_client_id,
+        path_priv_key,
+        logger=None,
+    ):
+        super().__init__(path_client_id, path_priv_key, logger)
+
+    def cdtools(self, cxiname, path_job_script, path_cdtools_nersc, n_gpu, **kwargs):
+        args_string = cdtool_args_string(
+            cxiname, path_cdtools_nersc, cdtools_parms, **kwargs
+        )
+        job_script = get_job_script(path_job_script, n_gpu, args_string)
+        self.logger.info(f"Job script: {job_script}")
+
+        self.submit_job(job_script)
+        self.task_wait()
+
+    def ptychocam(
+        self, cxiname, path_job_script, path_ptychocam_nersc, n_gpu, **kwargs
+    ):
+        args_string = ptychocam_args_string(
+            cxiname, path_ptychocam_nersc, ptychocam_parms, **kwargs
+        )
+        job_script = get_job_script(path_job_script, n_gpu, args_string)
+
+        self.logger.info(f"Job script: {job_script}")
+
+        self.submit_job(job_script)
+        self.task_wait()