Merge pull request #56 from neurostuff/enh/use_sdks

[ENH] use neurostore/neurosynth compose sdks

Author: jdkent

compose_runner/aws_lambda/common.py

@@ -53,8 +53,9 @@ def respond(self, body: Dict[str, Any], status_code: int = 200) -> Dict[str, Any
         return body
 
     def bad_request(self, message: str, status_code: int = 400) -> Dict[str, Any]:
-        return self.respond({"status": "FAILED", "error": message}, status_code=status_code)
+        return self.respond(
+            {"status": "FAILED", "error": message}, status_code=status_code
+        )
 
     def get(self, key: str, default: Any = None) -> Any:
         return self.payload.get(key, default)
-

compose_runner/aws_lambda/cost_check_handler.py

@@ -32,7 +32,11 @@ def _current_month_cost() -> Dict[str, Any]:
         Metrics=["UnblendedCost"],
     )
     results = response.get("ResultsByTime", [])
-    total = results[0]["Total"]["UnblendedCost"] if results else {"Amount": "0", "Unit": "USD"}
+    total = (
+        results[0]["Total"]["UnblendedCost"]
+        if results
+        else {"Amount": "0", "Unit": "USD"}
+    )
     amount = float(Decimal(total.get("Amount", "0")))
     currency = total.get("Unit", "USD")
     return {"amount": amount, "currency": currency, "time_period": period}

compose_runner/aws_lambda/log_poll_handler.py

@@ -8,7 +8,9 @@
 
 from compose_runner.aws_lambda.common import LambdaRequest
 
-_LOGS_CLIENT = boto3.client("logs", region_name=os.environ.get("AWS_REGION", "us-east-1"))
+_LOGS_CLIENT = boto3.client(
+    "logs", region_name=os.environ.get("AWS_REGION", "us-east-1")
+)
 
 LOG_GROUP_ENV = "RUNNER_LOG_GROUP"
 DEFAULT_LOOKBACK_MS_ENV = "DEFAULT_LOOKBACK_MS"

compose_runner/aws_lambda/results_handler.py

@@ -35,7 +35,9 @@ def handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]:
         raise KeyError(message)
     expires_in = int(payload.get("expires_in", DEFAULT_EXPIRES_IN))
 
-    key_prefix = f"{prefix.rstrip('/')}/{artifact_prefix}" if prefix else artifact_prefix
+    key_prefix = (
+        f"{prefix.rstrip('/')}/{artifact_prefix}" if prefix else artifact_prefix
+    )
 
     response = _S3.list_objects_v2(Bucket=bucket, Prefix=key_prefix)
     contents = response.get("Contents", [])

compose_runner/aws_lambda/run_handler.py

@@ -16,7 +16,9 @@
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
 
-_SFN_CLIENT = boto3.client("stepfunctions", region_name=os.environ.get("AWS_REGION", "us-east-1"))
+_SFN_CLIENT = boto3.client(
+    "stepfunctions", region_name=os.environ.get("AWS_REGION", "us-east-1")
+)
 
 STATE_MACHINE_ARN_ENV = "STATE_MACHINE_ARN"
 RESULTS_BUCKET_ENV = "RESULTS_BUCKET"
@@ -44,10 +46,14 @@ def _compose_api_base_url(environment: str) -> str:
     return "https://compose.neurosynth.org/api"
 
 
-def _fetch_meta_analysis(meta_analysis_id: str, environment: str) -> Optional[Dict[str, Any]]:
+def _fetch_meta_analysis(
+    meta_analysis_id: str, environment: str
+) -> Optional[Dict[str, Any]]:
     base_url = _compose_api_base_url(environment).rstrip("/")
     url = f"{base_url}/meta-analyses/{meta_analysis_id}?nested=true"
-    request = urllib.request.Request(url, headers={"User-Agent": "compose-runner/submit"})
+    request = urllib.request.Request(
+        url, headers={"User-Agent": "compose-runner/submit"}
+    )
     try:
         with urllib.request.urlopen(request, timeout=10) as response:
             return json.load(response)
@@ -77,7 +83,9 @@ def _requires_large_task(specification: Dict[str, Any]) -> bool:
     return False
 
 
-def _select_task_size(meta_analysis_id: str, environment: str, artifact_prefix: str) -> str:
+def _select_task_size(
+    meta_analysis_id: str, environment: str, artifact_prefix: str
+) -> str:
     doc = _fetch_meta_analysis(meta_analysis_id, environment)
     if not doc:
         return DEFAULT_TASK_SIZE
@@ -92,7 +100,9 @@ def _select_task_size(meta_analysis_id: str, environment: str, artifact_prefix:
             )
             return "large"
     except Exception as exc:  # noqa: broad-except
-        logger.warning("Failed to evaluate specification for %s: %s", meta_analysis_id, exc)
+        logger.warning(
+            "Failed to evaluate specification for %s: %s", meta_analysis_id, exc
+        )
     return DEFAULT_TASK_SIZE
 
 
@@ -146,9 +156,13 @@ def handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]:
     nv_key = payload.get("nv_key") or os.environ.get(NV_KEY_ENV)
 
     environment = payload.get("environment", "production")
-    task_size = _select_task_size(payload["meta_analysis_id"], environment, artifact_prefix)
+    task_size = _select_task_size(
+        payload["meta_analysis_id"], environment, artifact_prefix
+    )
 
-    job_input = _job_input(payload, artifact_prefix, bucket, prefix, nsc_key, nv_key, task_size)
+    job_input = _job_input(
+        payload, artifact_prefix, bucket, prefix, nsc_key, nv_key, task_size
+    )
     params = {
         "stateMachineArn": os.environ[STATE_MACHINE_ARN_ENV],
         "name": artifact_prefix,

compose_runner/aws_lambda/status_handler.py

@@ -10,7 +10,9 @@
 
 from compose_runner.aws_lambda.common import LambdaRequest
 
-_SFN = boto3.client("stepfunctions", region_name=os.environ.get("AWS_REGION", "us-east-1"))
+_SFN = boto3.client(
+    "stepfunctions", region_name=os.environ.get("AWS_REGION", "us-east-1")
+)
 _S3 = boto3.client("s3", region_name=os.environ.get("AWS_REGION", "us-east-1"))
 
 RESULTS_BUCKET_ENV = "RESULTS_BUCKET"
@@ -28,7 +30,9 @@ def _metadata_key(prefix: Optional[str], artifact_prefix: str) -> str:
     return f"{artifact_prefix}/{METADATA_FILENAME}"
 
 
-def _load_metadata(bucket: str, prefix: Optional[str], artifact_prefix: str) -> Optional[Dict[str, Any]]:
+def _load_metadata(
+    bucket: str, prefix: Optional[str], artifact_prefix: str
+) -> Optional[Dict[str, Any]]:
     key = _metadata_key(prefix, artifact_prefix)
     try:
         response = _S3.get_object(Bucket=bucket, Key=key)
@@ -65,7 +69,11 @@ def handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]:
     except ClientError as error:
         body = {"status": "FAILED", "error": error.response["Error"]["Message"]}
         if request.is_http:
-            status_code = 404 if error.response["Error"]["Code"] == "ExecutionDoesNotExist" else 500
+            status_code = (
+                404
+                if error.response["Error"]["Code"] == "ExecutionDoesNotExist"
+                else 500
+            )
             return request.respond(body, status_code=status_code)
         raise
 
@@ -83,7 +91,9 @@ def handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]:
 
     artifact_prefix = description.get("name")
     if not artifact_prefix:
-        raise ValueError("Execution does not expose a name; cannot determine artifact prefix.")
+        raise ValueError(
+            "Execution does not expose a name; cannot determine artifact prefix."
+        )
     body["artifact_prefix"] = artifact_prefix
 
     if status in {"SUCCEEDED", "FAILED"}:

compose_runner/cli.py

@@ -22,5 +22,7 @@ def cli(meta_analysis_id, environment, result_dir, nsc_key, nv_key, no_upload, n
 
     META_ANALYSIS_ID is the id of the meta-analysis on neurosynth-compose.
     """
-    url, _ = run(meta_analysis_id, environment, result_dir, nsc_key, nv_key, no_upload, n_cores)
+    url, _ = run(
+        meta_analysis_id, environment, result_dir, nsc_key, nv_key, no_upload, n_cores
+    )
     print(url)

compose_runner/ecs_task.py

@@ -48,15 +48,23 @@ def _iter_result_files(result_dir: Path) -> Iterable[Path]:
             yield path
 
 
-def _upload_results(artifact_prefix: str, result_dir: Path, bucket: str, prefix: Optional[str]) -> None:
-    base_prefix = f"{prefix.rstrip('/')}/{artifact_prefix}" if prefix else artifact_prefix
+def _upload_results(
+    artifact_prefix: str, result_dir: Path, bucket: str, prefix: Optional[str]
+) -> None:
+    base_prefix = (
+        f"{prefix.rstrip('/')}/{artifact_prefix}" if prefix else artifact_prefix
+    )
     for file_path in _iter_result_files(result_dir):
         key = f"{base_prefix}/{file_path.name}"
         _S3_CLIENT.upload_file(str(file_path), bucket, key)
 
 
-def _write_metadata(bucket: str, prefix: Optional[str], artifact_prefix: str, metadata: Dict[str, Any]) -> None:
-    base_prefix = f"{prefix.rstrip('/')}/{artifact_prefix}" if prefix else artifact_prefix
+def _write_metadata(
+    bucket: str, prefix: Optional[str], artifact_prefix: str, metadata: Dict[str, Any]
+) -> None:
+    base_prefix = (
+        f"{prefix.rstrip('/')}/{artifact_prefix}" if prefix else artifact_prefix
+    )
     key = f"{base_prefix}/{METADATA_FILENAME}"
     metadata["metadata_key"] = key
     _S3_CLIENT.put_object(