truss 0.10.0rc1__py3-none-any.whl → 0.60.0__py3-none-any.whl

truss/tests/remote/baseten/test_service.py

@@ -0,0 +1,55 @@
+from truss.remote.baseten import service
+
+
+def test_model_invocation_url_prod():
+    url = service.URLConfig.invocation_url(
+        "https://api.baseten.co", service.URLConfig.MODEL, "123", "789", is_draft=False
+    )
+    assert url == "https://model-123.api.baseten.co/deployment/789/predict"
+
+
+def test_model_invocation_url_draft():
+    url = service.URLConfig.invocation_url(
+        "https://api.baseten.co", service.URLConfig.MODEL, "123", "789", is_draft=True
+    )
+    assert url == "https://model-123.api.baseten.co/development/predict"
+
+
+def test_chain_invocation_url_prod():
+    url = service.URLConfig.invocation_url(
+        "https://api.baseten.co", service.URLConfig.CHAIN, "abc", "666", is_draft=False
+    )
+    assert url == "https://chain-abc.api.baseten.co/deployment/666/run_remote"
+
+
+def test_chain_invocation_url_draft():
+    url = service.URLConfig.invocation_url(
+        "https://api.baseten.co", service.URLConfig.CHAIN, "abc", "666", is_draft=True
+    )
+    assert url == "https://chain-abc.api.baseten.co/development/run_remote"
+
+
+def test_model_status_page_url():
+    url = service.URLConfig.status_page_url(
+        "https://app.baseten.co", service.URLConfig.MODEL, "123"
+    )
+    assert url == "https://app.baseten.co/models/123/overview"
+
+
+def test_chain_status_page_url():
+    url = service.URLConfig.status_page_url(
+        "https://app.baseten.co", service.URLConfig.CHAIN, "abc"
+    )
+    assert url == "https://app.baseten.co/chains/abc/overview"
+
+
+def test_model_logs_url():
+    url = service.URLConfig.model_logs_url("https://app.baseten.co", "123", "789")
+    assert url == "https://app.baseten.co/models/123/logs/789"
+
+
+def test_chain_logs_url():
+    url = service.URLConfig.chainlet_logs_url(
+        "https://app.baseten.co", "abc", "666", "543"
+    )
+    assert url == "https://app.baseten.co/chains/abc/logs/666/543"

truss/tests/remote/test_remote_factory.py

@@ -2,7 +2,7 @@ from unittest import mock
 
 import pytest
 from truss.remote.remote_factory import RemoteFactory
-from truss.remote.truss_remote import RemoteConfig, TrussRemote, TrussService
+from truss.remote.truss_remote import RemoteConfig, RemoteUser, TrussRemote
 
 SAMPLE_CONFIG = {"api_key": "test_key", "remote_url": "http://test.com"}
 
@@ -25,10 +25,9 @@ remote_provider=test_remote
 """
 
 
-class TestRemote(TrussRemote):
+class TrussTestRemote(TrussRemote):
     def __init__(self, api_key, remote_url):
         self.api_key = api_key
-        self.remote_url = remote_url
 
     def authenticate(self):
         return {"Authorization": self.api_key}
@@ -36,20 +35,19 @@ class TestRemote(TrussRemote):
     def push(self):
         return {"status": "success"}
 
-    def get_remote_logs_url(self, service: TrussService) -> str:
-        raise NotImplementedError
-
     def get_service(self, **kwargs):
         raise NotImplementedError
 
     def sync_truss_to_dev_version_by_name(self, model_name: str, target_directory: str):
         raise NotImplementedError
 
+    def whoami(self) -> RemoteUser:
+        return RemoteUser("test_user", "test_email")
+
 
 def mock_service_config():
     return RemoteConfig(
-        name="mock-service",
-        configs={"remote_provider": "test_remote", **SAMPLE_CONFIG},
+        name="mock-service", configs={"remote_provider": "test_remote", **SAMPLE_CONFIG}
     )
 
 
@@ -60,7 +58,7 @@ def mock_incorrect_service_config():
     )
 
 
-@mock.patch.dict(RemoteFactory.REGISTRY, {"test_remote": TestRemote}, clear=True)
+@mock.patch.dict(RemoteFactory.REGISTRY, {"test_remote": TrussTestRemote}, clear=True)
 @mock.patch(
     "truss.remote.remote_factory.RemoteFactory.load_remote_config",
     return_value=mock_service_config(),
@@ -69,10 +67,10 @@ def test_create(mock_load_remote_config):
     service_name = "test_service"
     remote = RemoteFactory.create(service_name)
     mock_load_remote_config.assert_called_once_with(service_name)
-    assert isinstance(remote, TestRemote)
+    assert isinstance(remote, TrussTestRemote)
 
 
-@mock.patch.dict(RemoteFactory.REGISTRY, {"test_remote": TestRemote}, clear=True)
+@mock.patch.dict(RemoteFactory.REGISTRY, {"test_remote": TrussTestRemote}, clear=True)
 @mock.patch(
     "truss.remote.remote_factory.RemoteFactory.load_remote_config",
     return_value=mock_incorrect_service_config(),
@@ -83,7 +81,7 @@ def test_create_no_service(mock_load_remote_config):
         RemoteFactory.create(service_name)
 
 
-@mock.patch.dict(RemoteFactory.REGISTRY, {"test_remote": TestRemote}, clear=True)
+@mock.patch.dict(RemoteFactory.REGISTRY, {"test_remote": TrussTestRemote}, clear=True)
 @mock.patch("builtins.open", new_callable=mock.mock_open, read_data=SAMPLE_TRUSSRC)
 @mock.patch("pathlib.Path.exists", return_value=True)
 def test_load_remote_config(mock_exists, mock_open):
@@ -92,7 +90,7 @@ def test_load_remote_config(mock_exists, mock_open):
     assert service.configs == {"remote_provider": "test_remote", **SAMPLE_CONFIG}
 
 
-@mock.patch.dict(RemoteFactory.REGISTRY, {"test_remote": TestRemote}, clear=True)
+@mock.patch.dict(RemoteFactory.REGISTRY, {"test_remote": TrussTestRemote}, clear=True)
 @mock.patch("builtins.open", new_callable=mock.mock_open, read_data=SAMPLE_TRUSSRC)
 @mock.patch("pathlib.Path.exists", return_value=False)
 def test_load_remote_config_no_file(mock_exists, mock_open):
@@ -100,7 +98,7 @@ def test_load_remote_config_no_file(mock_exists, mock_open):
         RemoteFactory.load_remote_config("test")
 
 
-@mock.patch.dict(RemoteFactory.REGISTRY, {"test_remote": TestRemote}, clear=True)
+@mock.patch.dict(RemoteFactory.REGISTRY, {"test_remote": TrussTestRemote}, clear=True)
 @mock.patch("builtins.open", new_callable=mock.mock_open, read_data=SAMPLE_TRUSSRC)
 @mock.patch("pathlib.Path.exists", return_value=True)
 def test_load_remote_config_no_service(mock_exists, mock_open):
@@ -108,13 +106,13 @@ def test_load_remote_config_no_service(mock_exists, mock_open):
         RemoteFactory.load_remote_config("nonexistent_service")
 
 
-@mock.patch.dict(RemoteFactory.REGISTRY, {"test_remote": TestRemote}, clear=True)
+@mock.patch.dict(RemoteFactory.REGISTRY, {"test_remote": TrussTestRemote}, clear=True)
 def test_required_params():
-    required_params = RemoteFactory.required_params(TestRemote)
+    required_params = RemoteFactory.required_params(TrussTestRemote)
     assert required_params == {"api_key", "remote_url"}
 
 
-@mock.patch.dict(RemoteFactory.REGISTRY, {"test_remote": TestRemote}, clear=True)
+@mock.patch.dict(RemoteFactory.REGISTRY, {"test_remote": TrussTestRemote}, clear=True)
 @mock.patch(
     "builtins.open", new_callable=mock.mock_open, read_data=SAMPLE_TRUSSRC_NO_REMOTE
 )
@@ -125,7 +123,7 @@ def test_validate_remote_config_no_remote(mock_exists, mock_open):
         RemoteFactory.validate_remote_config(service.configs, "test")
 
 
-@mock.patch.dict(RemoteFactory.REGISTRY, {"test_remote": TestRemote}, clear=True)
+@mock.patch.dict(RemoteFactory.REGISTRY, {"test_remote": TrussTestRemote}, clear=True)
 @mock.patch(
     "builtins.open", new_callable=mock.mock_open, read_data=SAMPLE_TRUSSRC_NO_PARAMS
 )

truss/tests/remote/test_truss_remote.py

@@ -1,3 +1,4 @@
+from typing import Iterator
 from unittest import mock
 
 import pytest
@@ -7,28 +8,36 @@ from truss.remote.truss_remote import TrussService
 TEST_SERVICE_URL = "http://test.com"
 
 
-class TestTrussService(TrussService):
-    def __init__(self, remote_url: str, is_draft: bool, **kwargs):
-        super().__init__(remote_url, is_draft, **kwargs)
-        self._remote_url = remote_url
+class TrussTestService(TrussService):
+    def __init__(self, _service_url: str, is_draft: bool, **kwargs):
+        super().__init__(_service_url, is_draft, **kwargs)
 
     def authenticate(self):
-        return {"Authorization": "Test"}
+        return {}
 
-    def is_live(self):
-        response = self._send_request(self._remote_url, "GET")
+    def is_live(self) -> bool:
+        response = self._send_request(self._service_url, "GET")
         if response.status_code == 200:
             return True
         return False
 
-    def is_ready(self):
-        response = self._send_request(self._remote_url, "GET")
+    def is_ready(self) -> bool:
+        response = self._send_request(self._service_url, "GET")
         if response.status_code == 200:
             return True
         return False
 
-    def logs_url(self, base_url: str) -> str:
-        return f"{base_url}/logs"
+    @property
+    def logs_url(self) -> str:
+        raise NotImplementedError()
+
+    @property
+    def predict_url(self) -> str:
+        return f"{self._service_url}/v1/models/model:predict"
+
+    def poll_deployment_status(self, sleep_secs: int = 1) -> Iterator[str]:
+        for status in ["DEPLOYING", "ACTIVE"]:
+            yield status
 
 
 def mock_successful_response():
@@ -46,37 +55,37 @@ def mock_unsuccessful_response():
 
 @mock.patch("requests.request", return_value=mock_successful_response())
 def test_is_live(mock_request):
-    service = TestTrussService(TEST_SERVICE_URL, True)
+    service = TrussTestService(TEST_SERVICE_URL, True)
     assert service.is_live()
 
 
 @mock.patch("requests.request", return_value=mock_unsuccessful_response())
 def test_is_not_live(mock_request):
-    service = TestTrussService(TEST_SERVICE_URL, True)
+    service = TrussTestService(TEST_SERVICE_URL, True)
     assert service.is_live() is False
 
 
 @mock.patch("requests.request", return_value=mock_successful_response())
 def test_is_ready(mock_request):
-    service = TestTrussService(TEST_SERVICE_URL, True)
+    service = TrussTestService(TEST_SERVICE_URL, True)
     assert service.is_ready()
 
 
 @mock.patch("requests.request", return_value=mock_unsuccessful_response())
 def test_is_not_ready(mock_request):
-    service = TestTrussService(TEST_SERVICE_URL, True)
+    service = TrussTestService(TEST_SERVICE_URL, True)
     assert service.is_ready() is False
 
 
 @mock.patch("requests.request", return_value=mock_successful_response())
 def test_predict(mock_request):
-    service = TestTrussService(TEST_SERVICE_URL, True)
+    service = TrussTestService(TEST_SERVICE_URL, True)
     response = service.predict({"model_input": "test"})
     assert response.status_code == 200
 
 
 @mock.patch("requests.request", return_value=mock_successful_response())
 def test_predict_no_data(mock_request):
-    service = TestTrussService(TEST_SERVICE_URL, True)
+    service = TrussTestService(TEST_SERVICE_URL, True)
     with pytest.raises(ValueError):
         service._send_request(TEST_SERVICE_URL, "POST")

truss/tests/templates/control/control/helpers/test_context_managers.py

@@ -0,0 +1,11 @@
+import os
+
+from truss.templates.control.control.helpers.context_managers import current_directory
+
+
+def test_current_directory(tmp_path):
+    orig_cwd = os.getcwd()
+    with current_directory(tmp_path):
+        assert os.getcwd() == str(tmp_path)
+
+    assert os.getcwd() == orig_cwd

truss/tests/templates/control/control/helpers/test_model_container_patch_applier.py

@@ -0,0 +1,184 @@
+import os
+import sys
+from pathlib import Path
+from unittest import mock
+
+import pytest
+from truss.base.truss_config import TrussConfig
+
+# Needed to simulate the set up on the model docker container
+sys.path.append(
+    str(
+        Path(__file__).parent.parent.parent.parent.parent.parent
+        / "templates"
+        / "control"
+        / "control"
+    )
+)
+
+# Have to use imports in this form, otherwise isinstance checks fail on helper classes
+from helpers.truss_patch.model_container_patch_applier import (  # noqa
+    ModelContainerPatchApplier,
+)
+from helpers.custom_types import (  # noqa
+    Action,
+    ConfigPatch,
+    EnvVarPatch,
+    ExternalDataPatch,
+    ModelCodePatch,
+    PackagePatch,
+    Patch,
+    PatchType,
+)
+
+
+@pytest.fixture
+def patch_applier(truss_container_fs):
+    return ModelContainerPatchApplier(truss_container_fs / "app", mock.Mock())
+
+
+def test_patch_applier_model_code_patch_add(
+    patch_applier: ModelContainerPatchApplier, truss_container_fs
+):
+    patch = Patch(
+        type=PatchType.MODEL_CODE,
+        body=ModelCodePatch(action=Action.ADD, path="dummy", content=""),
+    )
+    patch_applier(patch, os.environ.copy())
+    assert (truss_container_fs / "app" / "model" / "dummy").exists()
+
+
+def test_patch_applier_model_code_patch_remove(
+    patch_applier: ModelContainerPatchApplier, truss_container_fs
+):
+    patch = Patch(
+        type=PatchType.MODEL_CODE,
+        body=ModelCodePatch(action=Action.REMOVE, path="model.py"),
+    )
+    assert (truss_container_fs / "app" / "model" / "model.py").exists()
+    patch_applier(patch, os.environ.copy())
+    assert not (truss_container_fs / "app" / "model" / "model.py").exists()
+
+
+def test_patch_applier_model_code_patch_update(
+    patch_applier: ModelContainerPatchApplier, truss_container_fs
+):
+    new_model_file_content = """
+    class Model:
+        pass
+    """
+    patch = Patch(
+        type=PatchType.MODEL_CODE,
+        body=ModelCodePatch(
+            action=Action.UPDATE, path="model.py", content=new_model_file_content
+        ),
+    )
+    patch_applier(patch, os.environ.copy())
+    assert (
+        truss_container_fs / "app" / "model" / "model.py"
+    ).read_text() == new_model_file_content
+
+
+def test_patch_applier_package_patch_add(
+    patch_applier: ModelContainerPatchApplier, truss_container_fs
+):
+    patch = Patch(
+        type=PatchType.PACKAGE,
+        body=PackagePatch(
+            action=Action.ADD, path="test_package/test.py", content="foobar"
+        ),
+    )
+    patch_applier(patch, os.environ.copy())
+    assert (truss_container_fs / "packages" / "test_package" / "test.py").exists()
+
+
+def test_patch_applier_package_patch_remove(
+    patch_applier: ModelContainerPatchApplier, truss_container_fs
+):
+    patch = Patch(
+        type=PatchType.PACKAGE,
+        body=PackagePatch(action=Action.REMOVE, path="test_package/test.py"),
+    )
+    assert (truss_container_fs / "packages" / "test_package" / "test.py").exists()
+    patch_applier(patch, os.environ.copy())
+    assert not (truss_container_fs / "packages" / "test_package" / "test.py").exists()
+
+
+def test_patch_applier_package_patch_update(
+    patch_applier: ModelContainerPatchApplier, truss_container_fs
+):
+    new_package_content = """X = 2"""
+    patch = Patch(
+        type=PatchType.PACKAGE,
+        body=PackagePatch(
+            action=Action.UPDATE,
+            path="test_package/test.py",
+            content=new_package_content,
+        ),
+    )
+    patch_applier(patch, os.environ.copy())
+    assert (
+        truss_container_fs / "packages" / "test_package" / "test.py"
+    ).read_text() == new_package_content
+
+
+def test_patch_applier_config_patch_update(
+    patch_applier: ModelContainerPatchApplier, truss_container_fs
+):
+    new_config_dict = {"model_name": "foobar"}
+    patch = Patch(
+        type=PatchType.CONFIG,
+        body=ConfigPatch(action=Action.UPDATE, config=new_config_dict),
+    )
+    patch_applier(patch, os.environ.copy())
+    new_config = TrussConfig.from_yaml(truss_container_fs / "app" / "config.yaml")
+    assert new_config.model_name == "foobar"
+
+
+def test_patch_applier_env_var_patch_update(patch_applier: ModelContainerPatchApplier):
+    env_var_dict = {"FOO": "BAR"}
+    patch = Patch(
+        type=PatchType.ENVIRONMENT_VARIABLE,
+        body=EnvVarPatch(action=Action.UPDATE, item={"FOO": "BAR-PATCHED"}),
+    )
+    patch_applier(patch, env_var_dict)
+    assert env_var_dict["FOO"] == "BAR-PATCHED"
+
+
+def test_patch_applier_env_var_patch_add(patch_applier: ModelContainerPatchApplier):
+    env_var_dict = {"FOO": "BAR"}
+    patch = Patch(
+        type=PatchType.ENVIRONMENT_VARIABLE,
+        body=EnvVarPatch(action=Action.ADD, item={"BAR": "FOO"}),
+    )
+    patch_applier(patch, env_var_dict)
+    assert env_var_dict["FOO"] == "BAR"
+    assert env_var_dict["BAR"] == "FOO"
+
+
+def test_patch_applier_env_var_patch_remove(patch_applier: ModelContainerPatchApplier):
+    env_var_dict = {"FOO": "BAR"}
+    patch = Patch(
+        type=PatchType.ENVIRONMENT_VARIABLE,
+        body=EnvVarPatch(action=Action.REMOVE, item={"FOO": "BAR"}),
+    )
+    patch_applier(patch, env_var_dict)
+    with pytest.raises(KeyError):
+        _ = env_var_dict["FOO"]
+
+
+def test_patch_applier_external_data_patch_add(
+    patch_applier: ModelContainerPatchApplier, truss_container_fs
+):
+    patch = Patch(
+        type=PatchType.EXTERNAL_DATA,
+        body=ExternalDataPatch(
+            action=Action.ADD,
+            item={
+                "url": "https://raw.githubusercontent.com/basetenlabs/truss/main/docs/favicon.svg",
+                "local_data_path": "truss_icon",
+            },
+        ),
+    )
+    patch_applier(patch, os.environ.copy())
+    assert (truss_container_fs / "app" / "data" / "truss_icon").exists()

truss/tests/templates/control/control/helpers/test_requirement_name_identifier.py

@@ -0,0 +1,89 @@
+import pytest
+from truss.templates.control.control.helpers.truss_patch.requirement_name_identifier import (
+    RequirementMeta,
+    identify_requirement_name,
+    reqs_by_name,
+)
+
+
+@pytest.mark.parametrize(
+    "req, expected_name",
+    [
+        ("pytorch", "pytorch"),
+        (
+            "git+https://github.com/huggingface/transformers.git#egg=transformers",
+            "git+github.com/huggingface/transformers.git",
+        ),
+        (
+            "git+https://github.com/huggingface/transformers.git",
+            "git+github.com/huggingface/transformers.git",
+        ),
+        (
+            "git+https://github.com/huggingface/transformers.git@main#egg=transformers",
+            "git+github.com/huggingface/transformers.git",
+        ),
+        (
+            " git+https://github.com/huggingface/transformers.git ",
+            "git+github.com/huggingface/transformers.git",
+        ),
+        ("pytorch==1.0", "pytorch"),
+        ("pytorch>=1.0", "pytorch"),
+        ("pytorch<=1.0", "pytorch"),
+    ],
+)
+def test_identify_requirement_name(req, expected_name):
+    assert expected_name == identify_requirement_name(req)
+
+
+def test_reqs_by_name():
+    reqs = ["pytorch", " ", "jinja==1.0"]
+    assert reqs_by_name(reqs) == {"pytorch": "pytorch", "jinja": "jinja==1.0"}
+
+
+@pytest.mark.parametrize(
+    "desc, req, expected_meta",
+    [
+        (
+            "handles simple requirement",
+            "pytorch",
+            RequirementMeta(
+                requirement="pytorch",
+                name="pytorch",
+                is_url_based_requirement=False,
+                egg_tag=None,
+            ),
+        ),
+        (
+            "handles python package with version",
+            "pytorch==1.0",
+            RequirementMeta(
+                requirement="pytorch==1.0",
+                name="pytorch",
+                is_url_based_requirement=False,
+                egg_tag=None,
+            ),
+        ),
+        (
+            "handles url-based requirement with egg tag",
+            "git+https://github.com/huggingface/transformers.git@main#egg=transformers",
+            RequirementMeta(
+                requirement="git+https://github.com/huggingface/transformers.git@main#egg=transformers",
+                name="git+github.com/huggingface/transformers.git",
+                is_url_based_requirement=True,
+                egg_tag=["transformers"],
+            ),
+        ),
+        (
+            "handles url-based requirement without egg tag",
+            "git+https://github.com/huggingface/transformers.git",
+            RequirementMeta(
+                requirement="git+https://github.com/huggingface/transformers.git",
+                name="git+github.com/huggingface/transformers.git",
+                is_url_based_requirement=True,
+                egg_tag=None,
+            ),
+        ),
+    ],
+)
+def test_requirement_meta_from_req(desc, req: str, expected_meta: RequirementMeta):
+    assert expected_meta == RequirementMeta.from_req(req), desc