From 7500d23e58961003b144f9181ea5397cfca8f0e3 Mon Sep 17 00:00:00 2001
From: Krishna Kumar <krishnak@utexas.edu>
Date: Tue, 26 Nov 2024 15:42:37 -0600
Subject: [PATCH] MPM Job handler

---
 dapi/components/jobs/__init__.py             |  5 +-
 dapi/components/jobs/mpm_job_handler.py      | 48 ++++++++++++++++++++
 dapi/components/jobs/opensees_job_handler.py | 12 ++++-
 3 files changed, 63 insertions(+), 2 deletions(-)
 create mode 100644 dapi/components/jobs/mpm_job_handler.py

diff --git a/dapi/components/jobs/__init__.py b/dapi/components/jobs/__init__.py
index 4d344bc..8aafef2 100644
--- a/dapi/components/jobs/__init__.py
+++ b/dapi/components/jobs/__init__.py
@@ -16,9 +16,12 @@ def __init__(self, api):
 
     def _register_default_handlers(self) -> None:
         """Register default job handlers."""
-        from .opensees_job_handler import OpenSeesJobHandler  # Import default handlers
+        # Import default handlers
+        from .opensees_job_handler import OpenSeesJobHandler
+        from .mpm_job_handler import MpmJobHandler
 
         self.register_handler("opensees", OpenSeesJobHandler)
+        self.register_handler("mpm", MpmJobHandler)
 
     def register_handler(
         self, app_name: str, handler_class: Type[BaseJobHandler]
diff --git a/dapi/components/jobs/mpm_job_handler.py b/dapi/components/jobs/mpm_job_handler.py
new file mode 100644
index 0000000..bc7428f
--- /dev/null
+++ b/dapi/components/jobs/mpm_job_handler.py
@@ -0,0 +1,48 @@
+from typing import Dict, Any, Optional
+from .base_job_handler import BaseJobHandler
+
+
+class MpmJobHandler(BaseJobHandler):
+    """Custom handler for MPM (Material Point Method) jobs."""
+
+    def generate_job_info(
+        self,
+        tapis_client,
+        input_uri: str,
+        input_file: str,
+        job_name: Optional[str] = None,
+        max_minutes: Optional[int] = None,
+        node_count: Optional[int] = None,
+        cores_per_node: Optional[int] = None,
+        queue: Optional[str] = None,
+        allocation: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        # Get app info for the single MPM app
+        app_info = tapis_client.apps.getApp(appId="mpm")
+
+        # Create the base job info
+        job_info = {
+            "name": job_name,
+            "appId": app_info.id,
+            "appVersion": app_info.version,
+            "execSystemId": app_info.jobAttributes.execSystemId,
+            "maxMinutes": max_minutes or app_info.jobAttributes.maxMinutes,
+            "archiveOnAppError": app_info.jobAttributes.archiveOnAppError,
+            "fileInputs": [{"name": "Input Directory", "sourceUrl": input_uri}],
+            "execSystemLogicalQueue": queue
+            or app_info.jobAttributes.execSystemLogicalQueue,
+            "nodeCount": node_count or 1,  # Default to 1 if not specified
+            "coresPerNode": cores_per_node or 1,  # Default to 1 if not specified
+            "parameterSet": {
+                "appArgs": [{"name": "Input Script", "arg": input_file}],
+                "schedulerOptions": [],
+            },
+        }
+
+        # Add TACC allocation if provided
+        if allocation:
+            job_info["parameterSet"]["schedulerOptions"].append(
+                {"name": "TACC Allocation", "arg": f"-A {allocation}"}
+            )
+
+        return job_info
diff --git a/dapi/components/jobs/opensees_job_handler.py b/dapi/components/jobs/opensees_job_handler.py
index bd26da4..b731c2a 100644
--- a/dapi/components/jobs/opensees_job_handler.py
+++ b/dapi/components/jobs/opensees_job_handler.py
@@ -22,7 +22,17 @@ def generate_job_info(
         additional_params = {
             "envVariables": [{"key": "tclScript", "value": input_file}]
         }
-        app_name
+
+        # Use default app if no app specified or if specified app doesn't contain 'opensees'
+        if not app_name or "opensees" not in app_name.lower():
+            app_name = "opensees-express"
+        else:
+            # List available OpenSees apps
+            apps = tapis_client.apps.getApps(search="opensees")
+            if not any(app.id == app_name for app in apps):
+                raise ValueError(
+                    f"App '{app_name}' not found in available OpenSees apps"
+                )
 
         # Call the generic generate_job_info with OpenSees-specific params
         return super().generate_job_info(