Skip to content

Commit

Permalink
MPM Job handler
Browse files Browse the repository at this point in the history
  • Loading branch information
kks32 committed Nov 26, 2024
1 parent faee45f commit 7500d23
Show file tree
Hide file tree
Showing 3 changed files with 63 additions and 2 deletions.
5 changes: 4 additions & 1 deletion dapi/components/jobs/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down
48 changes: 48 additions & 0 deletions dapi/components/jobs/mpm_job_handler.py
Original file line number Diff line number Diff line change
@@ -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
12 changes: 11 additions & 1 deletion dapi/components/jobs/opensees_job_handler.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down

0 comments on commit 7500d23

Please sign in to comment.