Saving/Update Pipeline Services (#154)

* Saving/Update Pipeline Services

* Docstrings for new functions

* Pipeline creation unit test

---------

Co-authored-by: Thiago Castro Ferreira <[email protected]>
Co-authored-by: Thiago Castro Ferreira <[email protected]>

Author: 3 people

aixplain/factories/pipeline_factory.py

@@ -22,6 +22,7 @@
 """
 import json
 import logging
+import os
 from typing import Dict, List, Optional, Text, Union
 from aixplain.enums.data_type import DataType
 from aixplain.enums.function import Function
@@ -207,7 +208,7 @@ def list(
                 output_data_types = [output_data_types]
             payload["inputDataTypes"] = [data_type.value for data_type in output_data_types]
 
-        logging.info(f"Start service for POST List Dataset - {url} - {headers} - {json.dumps(payload)}")
+        logging.info(f"Start service for POST List Pipeline - {url} - {headers} - {json.dumps(payload)}")
         r = _request_with_retry("post", url, headers=headers, json=payload)
         resp = r.json()
 
@@ -220,3 +221,40 @@ def list(
             for pipeline in results:
                 pipelines.append(cls.__from_response(pipeline))
         return {"results": pipelines, "page_total": page_total, "page_number": page_number, "total": total}
+
+    @classmethod
+    def create(cls, name: Text, pipeline: Union[Text, Dict], status: Text = "draft") -> Pipeline:
+        """Pipeline Creation
+
+        Args:
+            name (Text): Pipeline Name
+            pipeline (Union[Text, Dict]): Pipeline as a Python dictionary or in a JSON file
+            status (Text, optional): Status of the pipeline. Currently only draft pipelines can be saved. Defaults to "draft".
+
+        Raises:
+            Exception: Currently just the creation of draft pipelines are supported
+
+        Returns:
+            Pipeline: instance of the new pipeline
+        """
+        try:
+            assert status == "draft", "Pipeline Creation Error: Currently just the creation of draft pipelines are supported."
+            if isinstance(pipeline, str) is True:
+                _, ext = os.path.splitext(pipeline)
+                assert (
+                    os.path.exists(pipeline) and ext == ".json"
+                ), "Pipeline Creation Error: Make sure the pipeline to be save is in a JSON file."
+                with open(pipeline) as f:
+                    pipeline = json.load(f)
+
+            # prepare payload
+            payload = {"name": name, "status": "draft", "architecture": pipeline}
+            url = urljoin(cls.backend_url, "sdk/pipelines")
+            headers = {"Authorization": f"Token {config.TEAM_API_KEY}", "Content-Type": "application/json"}
+            logging.info(f"Start service for POST Create Pipeline - {url} - {headers} - {json.dumps(payload)}")
+            r = _request_with_retry("post", url, headers=headers, json=payload)
+            response = r.json()
+
+            return Pipeline(response["id"], name, config.TEAM_API_KEY)
+        except Exception as e:
+            raise Exception(e)

aixplain/modules/pipeline.py

@@ -23,11 +23,13 @@
 
 import time
 import json
+import os
 import logging
 from aixplain.modules.asset import Asset
 from aixplain.utils import config
 from aixplain.utils.file_utils import _request_with_retry
 from typing import Dict, Optional, Text, Union
+from urllib.parse import urljoin
 
 
 class Pipeline(Asset):
@@ -306,3 +308,32 @@ def run_async(
             if resp is not None:
                 response["error"] = resp
         return response
+
+    def update(self, pipeline: Union[Text, Dict]):
+        """Update Pipeline
+
+        Args:
+            pipeline (Union[Text, Dict]): Pipeline as a Python dictionary or in a JSON file
+
+        Raises:
+            Exception: Make sure the pipeline to be save is in a JSON file.
+        """
+        try:
+            if isinstance(pipeline, str) is True:
+                _, ext = os.path.splitext(pipeline)
+                assert (
+                    os.path.exists(pipeline) and ext == ".json"
+                ), "Pipeline Update Error: Make sure the pipeline to be save is in a JSON file."
+                with open(pipeline) as f:
+                    pipeline = json.load(f)
+
+            # prepare payload
+            payload = {"name": self.name, "status": "draft", "architecture": pipeline}
+            url = urljoin(config.BACKEND_URL, f"sdk/pipelines/{self.id}")
+            headers = {"Authorization": f"Token {config.TEAM_API_KEY}", "Content-Type": "application/json"}
+            logging.info(f"Start service for PUT Update Pipeline - {url} - {headers} - {json.dumps(payload)}")
+            r = _request_with_retry("put", url, headers=headers, json=payload)
+            response = r.json()
+            logging.info(f"Pipeline {response['id']} Updated.")
+        except Exception as e:
+            raise Exception(e)

tests/functional/pipelines/create_test.py

@@ -0,0 +1,64 @@
+__author__ = "thiagocastroferreira"
+
+"""
+Copyright 2022 The aiXplain SDK authors
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import json
+import pytest
+from aixplain.factories import PipelineFactory
+from aixplain.modules import Pipeline
+from uuid import uuid4
+
+
+def test_create_pipeline_from_json():
+    pipeline_json = "tests/functional/pipelines/data/pipeline.json"
+    pipeline_name = str(uuid4())
+    pipeline = PipelineFactory.create(name=pipeline_name, pipeline=pipeline_json)
+
+    assert isinstance(pipeline, Pipeline)
+    assert pipeline.id != ""
+
+
+def test_create_pipeline_from_string():
+    pipeline_json = "tests/functional/pipelines/data/pipeline.json"
+    with open(pipeline_json) as f:
+        pipeline_dict = json.load(f)
+
+    pipeline_name = str(uuid4())
+    pipeline = PipelineFactory.create(name=pipeline_name, pipeline=pipeline_dict)
+
+    assert isinstance(pipeline, Pipeline)
+    assert pipeline.id != ""
+
+
+def test_update_pipeline():
+    pipeline_json = "tests/functional/pipelines/data/pipeline.json"
+    with open(pipeline_json) as f:
+        pipeline_dict = json.load(f)
+
+    pipeline_name = str(uuid4())
+    pipeline = PipelineFactory.create(name=pipeline_name, pipeline=pipeline_dict)
+
+    pipeline.update(pipeline=pipeline_json)
+    assert isinstance(pipeline, Pipeline)
+    assert pipeline.id != ""
+
+
+def test_create_pipeline_wrong_path():
+    pipeline_name = str(uuid4())
+
+    with pytest.raises(Exception):
+        pipeline = PipelineFactory.create(name=pipeline_name, pipeline="/")

tests/functional/pipelines/data/pipeline.json

@@ -0,0 +1,100 @@
+{
+    "links": [
+        {
+            "from": 0,
+            "to": 1,
+            "paramMapping": [
+                {
+                    "from": "input",
+                    "to": "text"
+                }
+            ]
+        },
+        {
+            "from": 1,
+            "to": 2,
+            "paramMapping": [
+                {
+                    "from": "data",
+                    "to": "text"
+                }
+            ]
+        },
+        {
+            "from": 2,
+            "to": 3,
+            "paramMapping": [
+                {
+                    "from": "data",
+                    "to": "output"
+                }
+            ]
+        }
+    ],
+    "nodes": [
+        {
+            "number": 0,
+            "type": "INPUT"
+        },
+        {
+            "number": 1,
+            "type": "ASSET",
+            "function": "sentiment-analysis",
+            "inputValues": [
+                {
+                    "code": "language",
+                    "value": "en"
+                },
+                {
+                    "code": "text",
+                    "dataType": "text"
+                }
+            ],
+            "assetId": "6172874f720b09325cbcdc33",
+            "assetType": "MODEL",
+            "autoSelectOptions": [],
+            "functionType": "AI",
+            "status": "Exists",
+            "outputValues": [
+                {
+                    "code": "data",
+                    "dataType": "label"
+                }
+            ]
+        },
+        {
+            "number": 2,
+            "type": "ASSET",
+            "function": "translation",
+            "inputValues": [
+                {
+                    "code": "sourcelanguage",
+                    "value": "en"
+                },
+                {
+                    "code": "targetlanguage",
+                    "value": "es"
+                },
+                {
+                    "code": "text",
+                    "dataType": "text"
+                }
+            ],
+            "assetId": "61b097551efecf30109d3316",
+            "assetType": "MODEL",
+            "autoSelectOptions": [],
+            "functionType": "AI",
+            "status": "Exists",
+            "outputValues": [
+                {
+                    "code": "data",
+                    "dataType": "text"
+                }
+            ]
+        },
+        {
+            "number": 3,
+            "type": "OUTPUT"
+        }
+    ]
+}

tests/unit/pipeline_test.py

@@ -0,0 +1,39 @@
+__author__ = "thiagocastroferreira"
+
+"""
+Copyright 2022 The aiXplain SDK authors
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+from dotenv import load_dotenv
+
+load_dotenv()
+import requests_mock
+from aixplain.utils import config
+from aixplain.factories import PipelineFactory
+from aixplain.modules import Pipeline
+from urllib.parse import urljoin
+import pytest
+
+
+def test_create_pipeline():
+    with requests_mock.Mocker() as mock:
+        url = urljoin(config.BACKEND_URL, "sdk/pipelines")
+        headers = {"x-api-key": config.TEAM_API_KEY, "Content-Type": "application/json"}
+        ref_response = {"id": "12345"}
+        mock.post(url, headers=headers, json=ref_response)
+        ref_pipeline = Pipeline(id="12345", name="Pipeline Test", api_key=config.TEAM_API_KEY)
+        hyp_pipeline = PipelineFactory.create(pipeline={}, name="Pipeline Test")
+    assert hyp_pipeline.id == ref_pipeline.id
+    assert hyp_pipeline.name == ref_pipeline.name