add a few tests for pipeline cloud (#72)

* add a few tests for pipeline cloud and have to cleanup too many flake8 errs

* Refactor schemes for future metrics changes (#74)

* Refactor schemes for future metrics changes

* Update version

* Add schemas for hardware related metrics (#75)

* Add schemas for hardware metrics

* Add project_id to metrics

* Refactor schemes for future metrics changes (#74)

* Refactor schemes for future metrics changes

* Update version

* Add schemas for hardware related metrics (#75)

* Add schemas for hardware metrics

* Add project_id to metrics

* add a few tests for pipeline cloud and have to cleanup too many flake8 errs

Co-authored-by: andrei-trandafir <47391556+andrei-trandafir@users.noreply.github.com>

examples/api_pipeline_run.py

@@ -1,6 +1,5 @@
 import json
 
-
 from pipeline import Pipeline, PipelineCloud, Variable, pipeline_function
 
 api = PipelineCloud()

examples/docker/basic_usage.py

@@ -1,5 +1,4 @@
-from pipeline import docker
-from pipeline import Pipeline, Variable, pipeline_function
+from pipeline import Pipeline, Variable, docker, pipeline_function
 
 
 @pipeline_function

examples/docker/hf_model.py

@@ -1,6 +1,6 @@
-from pipeline import docker
-from pipeline import Pipeline, Variable, pipeline_model, pipeline_function
-from transformers import GPTNeoForCausalLM, GPT2Tokenizer
+from transformers import GPT2Tokenizer, GPTNeoForCausalLM
+
+from pipeline import Pipeline, Variable, docker, pipeline_function, pipeline_model
 
 
 @pipeline_model
@@ -11,7 +11,7 @@ def __init__(self):
 
     @pipeline_function
     def predict(self, input_data: str) -> str:
-        if self.model == None:
+        if self.model is None:
             self.model = GPTNeoForCausalLM.from_pretrained("EleutherAI/gpt-neo-125M")
             self.tokenizer = GPT2Tokenizer.from_pretrained("EleutherAI/gpt-neo-125M")
 

examples/function_uploading.py

@@ -1,14 +1,10 @@
-import json
-
-
-from pipeline import Pipeline, PipelineCloud, Variable, pipeline_function
+from pipeline import PipelineCloud, pipeline_function
 
 api = PipelineCloud()
-api.authenticate()
 
 
 @pipeline_function
-def add_lol(a: str, b: float) -> str:
+def add_lol(a: str) -> str:
     return a + " lol"
 
 

examples/huggingface/gpt-j-6b.py

@@ -1,6 +1,10 @@
-from pipeline import Pipeline, PipelineCloud, Variable
-
-from pipeline import pipeline_model, pipeline_function
+from pipeline import (
+    Pipeline,
+    PipelineCloud,
+    Variable,
+    pipeline_function,
+    pipeline_model,
+)
 
 
 @pipeline_model
@@ -43,7 +47,10 @@ def predict(self, input_data: str, model_kwargs: dict = {}, **kwargs) -> str:
             )
         ):
             return {
-                "error": "GPT-J inference is limited to 2048 tokens. Reduce the prompt length and/or the expected generation length."
+                "error": (
+                    "GPT-J inference is limited to 2048 tokens.",
+                    "Reduce the prompt length and/or the expected generation length.",
+                )
             }
         if "remove_input" not in kwargs:
             kwargs["remove_input"] = False
@@ -79,9 +86,9 @@ def predict(self, input_data: str, model_kwargs: dict = {}, **kwargs) -> str:
 
     @pipeline_function
     def load(self) -> None:
-        from transformers import AutoModelForCausalLM, AutoTokenizer
-        import transformers
         import torch
+        import transformers
+        from transformers import AutoModelForCausalLM, AutoTokenizer
 
         class no_init:
             def __init__(self, modules=None, use_hf_no_init=True):

examples/huggingface/gpt-neo-upload.py

@@ -3,10 +3,13 @@
 )"""
 
 
-from pipeline import Pipeline, PipelineCloud, Variable
-
-from pipeline import pipeline_model, pipeline_function
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from pipeline import (
+    Pipeline,
+    PipelineCloud,
+    Variable,
+    pipeline_function,
+    pipeline_model,
+)
 
 
 @pipeline_model
@@ -31,7 +34,6 @@ def predict(self, input_data: str, model_kwargs: dict = {}) -> str:
     @pipeline_function
     def load(self) -> None:
         from transformers import AutoModelForCausalLM, AutoTokenizer
-        import transformers
 
         if self.model is None:
             self.model = AutoModelForCausalLM.from_pretrained(self.model_path)

pipeline/objects/huggingface/TransformersModelForCausalLM.py

@@ -1,6 +1,7 @@
-from pipeline import pipeline_model, pipeline_function
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
+from pipeline import pipeline_function, pipeline_model
+
 
 @pipeline_model()
 class TransformersModelForCausalLM:

pipeline/objects/model.py

@@ -2,9 +2,8 @@
 from hashlib import sha256
 from typing import Any
 
-from pipeline.util import generate_id, hex_to_python_object
-
 from pipeline.schemas.model import ModelGet
+from pipeline.util import generate_id, hex_to_python_object
 
 
 class Model:

pipeline/schemas/run.py

@@ -1,6 +1,6 @@
 import datetime
 from enum import Enum
-from typing import List, Optional, Union, Any
+from typing import Any, List, Optional, Union
 
 from pydantic import root_validator
 

pyproject.toml

@@ -37,6 +37,7 @@ pre-commit = "^2.16.0"
 flake8 = "^4.0.1"
 isort = "^5.10.1"
 numpy = "^1.22.0"
+responses = "^0.20.0"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]

setup.cfg

@@ -1,4 +1,4 @@
 [flake8]
 max-line-length = 88
 extend-ignore = "E203"
-per-file-ignores = tests/schemas/test_imports.py:F401
+per-file-ignores = tests/schemas/test_imports.py:F401 tests/test_serialization.py:E501 examples/huggingface/gpt-j-6b.py:E203

tests/conftest.py

@@ -0,0 +1,88 @@
+from datetime import datetime
+
+import pytest
+import responses
+from responses import matchers
+
+from pipeline.schemas.data import DataGet
+from pipeline.schemas.file import FileGet
+
+python_content = """
+from pipeline.objects import Pipeline, Variable, pipeline_function
+
+
+# Check if the decorator correctly uses __init__ and __enter__
+def test_with_decorator():
+    with Pipeline("test"):
+        assert Pipeline._current_pipeline is not None
+"""
+
+
+@pytest.fixture
+def api_response(url, token, bad_token, file_get_json):
+    with responses.RequestsMock(assert_all_requests_are_fired=False) as rsps:
+        rsps.add(
+            responses.GET,
+            url + "/v2/users/me",
+            json={"auth": True},
+            status=200,
+            match=[matchers.header_matcher({"Authorization": "Bearer " + token})],
+        )
+        rsps.add(
+            responses.GET,
+            url + "/v2/users/me",
+            json={"auth": True},
+            status=401,
+            match=[matchers.header_matcher({"Authorization": "Bearer " + bad_token})],
+        )
+        rsps.add(
+            responses.POST,
+            url + "/v2/files/",
+            json=file_get_json,
+            status=201,
+            match=[matchers.header_matcher({"Authorization": "Bearer " + token})],
+        )
+        yield rsps
+
+
+@pytest.fixture()
+def url():
+    return "http://127.0.0.1:8080"
+
+
+@pytest.fixture()
+def token():
+    return "token"
+
+
+@pytest.fixture()
+def bad_token():
+    return "bad_token"
+
+
+@pytest.fixture()
+def tmp_file():
+    return "tests/test_model.py"
+
+
+@pytest.fixture()
+def file_get_json():
+    return {
+        "name": "test",
+        "id": "file_test",
+        "path": "test/path/to/file",
+        "data": "data",
+        "file_size": 8,
+    }
+
+
+@pytest.fixture()
+def file_get():
+    return FileGet(
+        name="test", id="file_test", path="test/path/to/file", data="data", file_size=8
+    )
+
+
+@pytest.fixture()
+def data_get(file_get):
+    return DataGet(id="data_test", hex_file=file_get, created_at=datetime.now())

tests/test_pipeline_cloud.py

@@ -0,0 +1,25 @@
+import http
+
+import pytest
+
+from pipeline import PipelineCloud
+
+
+@pytest.mark.usefixtures("api_response")
+def test_cloud_init(url, token):
+    api = PipelineCloud(url, token)
+    assert api.token == token
+
+
+@pytest.mark.usefixtures("api_response")
+def test_cloud_init_failure(url, bad_token):
+    with pytest.raises(Exception) as e:
+        PipelineCloud(url, bad_token)
+        assert e.status == http.HTTPStatus.UNAUTHORIZED
+
+
+@pytest.mark.usefixtures("api_response")
+def test_cloud_upload_file(url, token, file_get, tmp_file):
+    api = PipelineCloud(url, token)
+    f = api.upload_file(tmp_file, "remote_path")
+    assert f == file_get

tests/test_serialization.py

@@ -1,5 +1,4 @@
 from pipeline.objects.graph import Graph
-
 from pipeline.schemas.pipeline import PipelineGet
 
 test_pipeline_dict = {