Support running analysis from agent via SDK (#126)

* support running analysis from agent

* lint

Author: nathanthorpe

cirro/cirro_client.py

@@ -4,7 +4,7 @@
 from cirro.auth.base import AuthInfo
 from cirro.config import AppConfig
 from cirro.services import FileService, DatasetService, ProjectService, ProcessService, ExecutionService, \
-    MetricsService, MetadataService, BillingService, ReferenceService, UserService
+    MetricsService, MetadataService, BillingService, ReferenceService, UserService, ComputeEnvironmentService
 
 
 class CirroApi:
@@ -51,6 +51,7 @@ def __init__(self, auth_info: AuthInfo = None, base_url: str = None):
         self._project_service = ProjectService(self._api_client)
         self._process_service = ProcessService(self._api_client)
         self._execution_service = ExecutionService(self._api_client)
+        self._compute_environment_service = ComputeEnvironmentService(self._api_client)
         self._metrics_service = MetricsService(self._api_client)
         self._metadata_service = MetadataService(self._api_client)
         self._billing_service = BillingService(self._api_client)
@@ -85,6 +86,13 @@ def execution(self) -> ExecutionService:
         """
         return self._execution_service
 
+    @property
+    def compute_environments(self) -> ComputeEnvironmentService:
+        """
+        List and update compute environments
+        """
+        return self._compute_environment_service
+
     @property
     def metrics(self) -> MetricsService:
         """

cirro/sdk/dataset.py

@@ -173,7 +173,8 @@ def run_analysis(
             description: str = "",
             process: Union[DataPortalProcess, str] = None,
             params=None,
-            notifications_emails=None
+            notifications_emails=None,
+            compute_environment=None
     ) -> str:
         """
         Runs an analysis on a dataset, returns the ID of the newly created dataset.
@@ -187,6 +188,8 @@ def run_analysis(
             process (DataPortalProcess or str): Process to run
             params (dict): Analysis parameters
             notifications_emails (List[str]): Notification email address(es)
+            compute_environment (str): Name or ID of compute environment to use,
+             if blank it will run in AWS
 
         Returns:
             dataset_id (str): ID of newly created dataset
@@ -203,6 +206,18 @@ def run_analysis(
         # If the process is a string, try to parse it as a process name or ID
         process = parse_process_name_or_id(process, self._client)
 
+        if compute_environment:
+            compute_environments = self._client.compute_environments.list_environments_for_project(
+                project_id=self.project_id
+            )
+            compute_environment = next(
+                (env for env in compute_environments
+                 if env.name == compute_environment or env.id == compute_environment),
+                None
+            )
+            if compute_environment is None:
+                raise DataPortalInputError(f"Compute environment '{compute_environment}' not found")
+
         resp = self._client.execution.run_analysis(
             project_id=self.project_id,
             request=RunAnalysisRequest(
@@ -211,7 +226,8 @@ def run_analysis(
                 process_id=process.id,
                 source_dataset_ids=[self.id],
                 params=RunAnalysisRequestParams.from_dict(params),
-                notification_emails=notifications_emails
+                notification_emails=notifications_emails,
+                compute_environment_id=compute_environment.id if compute_environment else None
             )
         )
         return resp.id

cirro/services/__init__.py

@@ -1,4 +1,5 @@
 from .billing import BillingService
+from .compute_environment import ComputeEnvironmentService
 from .dataset import DatasetService
 from .execution import ExecutionService
 from .file import FileService
@@ -13,6 +14,7 @@
     'BillingService',
     'DatasetService',
     'ExecutionService',
+    'ComputeEnvironmentService',
     'FileService',
     'MetadataService',
     'MetricsService',

cirro/services/compute_environment.py

@@ -0,0 +1,25 @@
+from typing import List
+
+from cirro_api_client.v1.api.compute_environment import get_compute_environments
+from cirro_api_client.v1.models import ComputeEnvironmentConfiguration
+
+from cirro.services.base import BaseService
+
+
+class ComputeEnvironmentService(BaseService):
+    """
+    Service for interacting with the Compute Environment endpoints
+    """
+
+    def list_environments_for_project(self, project_id: str) -> List[ComputeEnvironmentConfiguration]:
+        """
+        List of custom compute environments for a project (i.e., an agent)
+
+        Args:
+            project_id (str): Project ID
+
+        Returns:
+            List of compute environments that are available for the project
+        """
+        return get_compute_environments.sync(project_id=project_id,
+                                             client=self._api_client)