truss 0.11.18rc500__py3-none-any.whl → 0.11.24rc2__py3-none-any.whl

truss/tests/remote/baseten/test_remote.py

@@ -1,3 +1,5 @@
+from unittest import mock
+
 import pydantic
 import pytest
 import requests_mock
@@ -11,16 +13,17 @@ from truss.remote.baseten.core import (
 )
 from truss.remote.baseten.custom_types import ChainletDataAtomic, OracleData
 from truss.remote.baseten.error import RemoteError
-from truss.remote.baseten.remote import BasetenRemote
 from truss.truss_handle.truss_handle import TrussHandle
 
 _TEST_REMOTE_URL = "http://test_remote.com"
 _TEST_REMOTE_GRAPHQL_PATH = "http://test_remote.com/graphql/"
 
-
-@pytest.fixture
-def remote():
-    return BasetenRemote(_TEST_REMOTE_URL, "api_key")
+TRUSS_RC_CONTENT = """
+[baseten]
+remote_provider = baseten
+api_key = test_key
+remote_url = http://test.com
+""".strip()
 
 
 def assert_request_matches_expected_query(request, expected_query) -> None:
@@ -269,30 +272,41 @@ def test_push_raised_value_error_when_keep_previous_prod_settings_and_not_promot
             )
 
 
-def test_create_chain_with_no_publish(remote):
-    with requests_mock.Mocker() as m:
-        m.post(
-            _TEST_REMOTE_GRAPHQL_PATH,
-            [
-                {"json": {"data": {"chains": []}}},
-                {
-                    "json": {
-                        "data": {
-                            "deploy_chain_atomic": {
-                                "chain_deployment": {
-                                    "id": "new-chain-deployment-id",
-                                    "chain": {
-                                        "id": "new-chain-id",
-                                        "hostname": "hostname",
-                                    },
-                                }
-                            }
-                        }
-                    }
-                },
-            ],
+@pytest.mark.parametrize("deploy_timeout_minutes", [9, 1441])
+def test_push_raised_value_error_when_deploy_timeout_minutes_is_invalid(
+    deploy_timeout_minutes, custom_model_truss_dir_with_pre_and_post, remote
+):
+    th = TrussHandle(custom_model_truss_dir_with_pre_and_post)
+
+    with pytest.raises(
+        ValueError,
+        match="deploy-timeout-minutes must be between 10 minutes and 1440 minutes \(24 hours\)",
+    ):
+        remote.push(
+            th,
+            "model_name",
+            th.truss_dir,
+            publish=True,
+            promote=False,
+            preserve_previous_prod_deployment=False,
+            deployment_name="dep_name",
+            deploy_timeout_minutes=deploy_timeout_minutes,
         )
 
+
+def test_create_chain_with_no_publish(remote):
+    mock_deploy_response = {
+        "chain_deployment": {
+            "id": "new-chain-deployment-id",
+            "chain": {"id": "new-chain-id", "hostname": "hostname"},
+        }
+    }
+
+    with mock.patch.object(
+        remote.api, "get_chains", return_value=[]
+    ) as mock_get_chains, mock.patch.object(
+        remote.api, "deploy_chain_atomic", return_value=mock_deploy_response
+    ) as mock_deploy:
         deployment_handle = create_chain_atomic(
             api=remote.api,
             chain_name="draft_chain",
@@ -310,64 +324,61 @@ def test_create_chain_with_no_publish(remote):
             environment=None,
         )
 
-        get_chains_graphql_request = m.request_history[0]
-        create_chain_graphql_request = m.request_history[1]
+        mock_get_chains.assert_called_once()
+        mock_deploy.assert_called_once()
 
-        expected_get_chains_query = """
-            {
-                chains {
-                    id
-                    name
-                }
-            }
-        """.strip()
+        call_kwargs = mock_deploy.call_args.kwargs
+        assert call_kwargs["chain_name"] == "draft_chain"
+        assert call_kwargs.get("is_draft") is True
+        assert call_kwargs.get("deploy_timeout_minutes") is None
 
-        assert_request_matches_expected_query(
-            get_chains_graphql_request, expected_get_chains_query
-        )
+        assert deployment_handle.chain_id == "new-chain-id"
+        assert deployment_handle.chain_deployment_id == "new-chain-deployment-id"
 
-        chainlets_string = """
-            {
-                name: "chainlet-1",
-                oracle: {
-                    model_name: "model-1",
-                    s3_key: "s3-key-1",
-                    encoded_config_str: "encoded-config-str-1",
-                    semver_bump: "MINOR"
-                }
-            }
-        """.strip()
-
-        # Note that if publish=False and promote=True, we set publish to True and create
-        # a non-draft deployment
-        expected_create_chain_mutation = f"""
-            mutation ($trussUserEnv: String) {{
-                deploy_chain_atomic(
-                    chain_name: "draft_chain"
-                    is_draft: true
-                    entrypoint: {chainlets_string}
-                    dependencies: []
-                    truss_user_env: $trussUserEnv
-                ) {{
-                    chain_deployment {{
-                        id
-                        chain {{
-                            id
-                            hostname
-                        }}
-                    }}
-                }}
-            }}
-        """.strip()
-
-        assert_request_matches_expected_query(
-            create_chain_graphql_request, expected_create_chain_mutation
+
+def test_create_chain_no_existing_chain(remote):
+    mock_deploy_response = {
+        "chain_deployment": {
+            "id": "new-chain-deployment-id",
+            "chain": {"id": "new-chain-id", "hostname": "hostname"},
+        }
+    }
+
+    with mock.patch.object(
+        remote.api, "get_chains", return_value=[]
+    ) as mock_get_chains, mock.patch.object(
+        remote.api, "deploy_chain_atomic", return_value=mock_deploy_response
+    ) as mock_deploy:
+        deployment_handle = create_chain_atomic(
+            api=remote.api,
+            chain_name="new_chain",
+            entrypoint=ChainletDataAtomic(
+                name="chainlet-1",
+                oracle=OracleData(
+                    model_name="model-1",
+                    s3_key="s3-key-1",
+                    encoded_config_str="encoded-config-str-1",
+                ),
+            ),
+            dependencies=[],
+            truss_user_env=b10_types.TrussUserEnv.collect(),
+            is_draft=False,
+            environment=None,
         )
+
+        mock_get_chains.assert_called_once()
+        mock_deploy.assert_called_once()
+
+        call_kwargs = mock_deploy.call_args.kwargs
+        assert call_kwargs["chain_name"] == "new_chain"
+        assert call_kwargs.get("is_draft") is not True
+        assert call_kwargs.get("deploy_timeout_minutes") is None
+
         assert deployment_handle.chain_id == "new-chain-id"
         assert deployment_handle.chain_deployment_id == "new-chain-deployment-id"
 
 
-def test_create_chain_no_existing_chain(remote):
+def test_create_chain_with_deployment_name(remote):
     with requests_mock.Mocker() as m:
         m.post(
             _TEST_REMOTE_GRAPHQL_PATH,
@@ -391,7 +402,8 @@ def test_create_chain_no_existing_chain(remote):
             ],
         )
 
-        deployment_handle = create_chain_atomic(
+        deployment_name = "chain-deployment"
+        create_chain_atomic(
             api=remote.api,
             chain_name="new_chain",
             entrypoint=ChainletDataAtomic(
@@ -406,94 +418,32 @@ def test_create_chain_no_existing_chain(remote):
             truss_user_env=b10_types.TrussUserEnv.collect(),
             is_draft=False,
             environment=None,
+            deployment_name=deployment_name,
         )
 
-        get_chains_graphql_request = m.request_history[0]
         create_chain_graphql_request = m.request_history[1]
 
-        expected_get_chains_query = """
-            {
-                chains {
-                    id
-                    name
-                }
-            }
-        """.strip()
-
-        assert_request_matches_expected_query(
-            get_chains_graphql_request, expected_get_chains_query
-        )
-
-        chainlets_string = """
-            {
-                name: "chainlet-1",
-                oracle: {
-                    model_name: "model-1",
-                    s3_key: "s3-key-1",
-                    encoded_config_str: "encoded-config-str-1",
-                    semver_bump: "MINOR"
-                }
-            }
-        """.strip()
-
-        expected_create_chain_mutation = f"""
-            mutation ($trussUserEnv: String) {{
-                deploy_chain_atomic(
-                    chain_name: "new_chain"
-                    is_draft: false
-                    entrypoint: {chainlets_string}
-                    dependencies: []
-                    truss_user_env: $trussUserEnv
-                ) {{
-                    chain_deployment {{
-                        id
-                        chain {{
-                            id
-                            hostname
-                        }}
-                    }}
-                }}
-            }}
-        """.strip()
-
-        assert_request_matches_expected_query(
-            create_chain_graphql_request, expected_create_chain_mutation
+        assert (
+            'deployment_name: "chain-deployment"'
+            in create_chain_graphql_request.json()["query"]
         )
 
-        assert deployment_handle.chain_id == "new-chain-id"
-        assert deployment_handle.chain_deployment_id == "new-chain-deployment-id"
-
 
 def test_create_chain_with_existing_chain_promote_to_environment_publish_false(remote):
-    with requests_mock.Mocker() as m:
-        m.post(
-            _TEST_REMOTE_GRAPHQL_PATH,
-            [
-                {
-                    "json": {
-                        "data": {
-                            "chains": [{"id": "old-chain-id", "name": "old_chain"}]
-                        }
-                    }
-                },
-                {
-                    "json": {
-                        "data": {
-                            "deploy_chain_atomic": {
-                                "chain_deployment": {
-                                    "id": "new-chain-deployment-id",
-                                    "chain": {
-                                        "id": "new-chain-id",
-                                        "hostname": "hostname",
-                                    },
-                                }
-                            }
-                        }
-                    }
-                },
-            ],
-        )
+    mock_deploy_response = {
+        "chain_deployment": {
+            "id": "new-chain-deployment-id",
+            "chain": {"id": "new-chain-id", "hostname": "hostname"},
+        }
+    }
 
+    with mock.patch.object(
+        remote.api,
+        "get_chains",
+        return_value=[{"id": "old-chain-id", "name": "old_chain"}],
+    ) as mock_get_chains, mock.patch.object(
+        remote.api, "deploy_chain_atomic", return_value=mock_deploy_response
+    ) as mock_deploy:
         deployment_handle = create_chain_atomic(
             api=remote.api,
             chain_name="old_chain",
@@ -511,95 +461,34 @@ def test_create_chain_with_existing_chain_promote_to_environment_publish_false(r
             environment="production",
         )
 
-        get_chains_graphql_request = m.request_history[0]
-        create_chain_graphql_request = m.request_history[1]
+        mock_get_chains.assert_called_once()
+        mock_deploy.assert_called_once()
 
-        expected_get_chains_query = """
-            {
-                chains {
-                    id
-                    name
-                }
-            }
-        """.strip()
-
-        assert_request_matches_expected_query(
-            get_chains_graphql_request, expected_get_chains_query
-        )
-
-        # Note that if publish=False and environment!=None, we set publish to True and create
-        # a non-draft deployment
-        chainlets_string = """
-            {
-                name: "chainlet-1",
-                oracle: {
-                    model_name: "model-1",
-                    s3_key: "s3-key-1",
-                    encoded_config_str: "encoded-config-str-1",
-                    semver_bump: "MINOR"
-                }
-            }
-        """.strip()
-
-        expected_create_chain_mutation = f"""
-            mutation ($trussUserEnv: String) {{
-                deploy_chain_atomic(
-                    chain_id: "old-chain-id"
-                    environment: "production"
-                    is_draft: false
-                    entrypoint: {chainlets_string}
-                    dependencies: []
-                    truss_user_env: $trussUserEnv
-                ) {{
-                    chain_deployment {{
-                        id
-                        chain {{
-                            id
-                            hostname
-                        }}
-                    }}
-                }}
-            }}
-        """.strip()
-
-        assert_request_matches_expected_query(
-            create_chain_graphql_request, expected_create_chain_mutation
-        )
+        call_kwargs = mock_deploy.call_args.kwargs
+        assert call_kwargs["chain_id"] == "old-chain-id"
+        assert call_kwargs["environment"] == "production"
+        assert call_kwargs.get("is_draft") is not True
+        assert call_kwargs.get("deploy_timeout_minutes") is None
 
         assert deployment_handle.chain_id == "new-chain-id"
         assert deployment_handle.chain_deployment_id == "new-chain-deployment-id"
 
 
 def test_create_chain_existing_chain_publish_true_no_promotion(remote):
-    with requests_mock.Mocker() as m:
-        m.post(
-            _TEST_REMOTE_GRAPHQL_PATH,
-            [
-                {
-                    "json": {
-                        "data": {
-                            "chains": [{"id": "old-chain-id", "name": "old_chain"}]
-                        }
-                    }
-                },
-                {
-                    "json": {
-                        "data": {
-                            "deploy_chain_atomic": {
-                                "chain_deployment": {
-                                    "id": "new-chain-deployment-id",
-                                    "chain": {
-                                        "id": "new-chain-id",
-                                        "hostname": "hostname",
-                                    },
-                                }
-                            }
-                        }
-                    }
-                },
-            ],
-        )
+    mock_deploy_response = {
+        "chain_deployment": {
+            "id": "new-chain-deployment-id",
+            "chain": {"id": "new-chain-id", "hostname": "hostname"},
+        }
+    }
 
+    with mock.patch.object(
+        remote.api,
+        "get_chains",
+        return_value=[{"id": "old-chain-id", "name": "old_chain"}],
+    ) as mock_get_chains, mock.patch.object(
+        remote.api, "deploy_chain_atomic", return_value=mock_deploy_response
+    ) as mock_deploy:
         deployment_handle = create_chain_atomic(
             api=remote.api,
             chain_name="old_chain",
@@ -617,57 +506,13 @@ def test_create_chain_existing_chain_publish_true_no_promotion(remote):
             environment=None,
         )
 
-        get_chains_graphql_request = m.request_history[0]
-        create_chain_graphql_request = m.request_history[1]
-
-        expected_get_chains_query = """
-            {
-                chains {
-                    id
-                    name
-                }
-            }
-        """.strip()
-
-        assert_request_matches_expected_query(
-            get_chains_graphql_request, expected_get_chains_query
-        )
+        mock_get_chains.assert_called_once()
+        mock_deploy.assert_called_once()
 
-        chainlets_string = """
-            {
-                name: "chainlet-1",
-                oracle: {
-                    model_name: "model-1",
-                    s3_key: "s3-key-1",
-                    encoded_config_str: "encoded-config-str-1",
-                    semver_bump: "MINOR"
-                }
-            }
-        """.strip()
-
-        expected_create_chain_mutation = f"""
-            mutation ($trussUserEnv: String) {{
-                deploy_chain_atomic(
-                    chain_id: "old-chain-id"
-                    is_draft: false
-                    entrypoint: {chainlets_string}
-                    dependencies: []
-                    truss_user_env: $trussUserEnv
-                ) {{
-                    chain_deployment {{
-                        id
-                        chain {{
-                            id
-                            hostname
-                        }}
-                    }}
-                }}
-            }}
-        """.strip()
-
-        assert_request_matches_expected_query(
-            create_chain_graphql_request, expected_create_chain_mutation
-        )
+        call_kwargs = mock_deploy.call_args.kwargs
+        assert call_kwargs["chain_id"] == "old-chain-id"
+        assert call_kwargs.get("is_draft") is not True
+        assert call_kwargs.get("deploy_timeout_minutes") is None
 
         assert deployment_handle.chain_id == "new-chain-id"
         assert deployment_handle.chain_deployment_id == "new-chain-deployment-id"
@@ -726,3 +571,86 @@ def test_push_raised_validation_error_for_extra_fields(tmp_path, remote):
             match="Extra fields not allowed: \[extra_field, who_am_i\]",
         ):
             remote.push(th, "model_name", th.truss_dir)
+
+
+def test_push_passes_deploy_timeout_minutes_to_create_truss_service(
+    custom_model_truss_dir_with_pre_and_post,
+    remote,
+    mock_baseten_requests,
+    mock_upload_truss,
+    mock_create_truss_service,
+    mock_truss_handle,
+):
+    remote.push(
+        mock_truss_handle,
+        "model_name",
+        mock_truss_handle.truss_dir,
+        publish=True,
+        deploy_timeout_minutes=450,
+    )
+
+    mock_create_truss_service.assert_called_once()
+    _, kwargs = mock_create_truss_service.call_args
+    assert kwargs["deploy_timeout_minutes"] == 450
+
+
+def test_push_passes_none_deploy_timeout_minutes_when_not_specified(
+    custom_model_truss_dir_with_pre_and_post,
+    remote,
+    mock_baseten_requests,
+    mock_upload_truss,
+    mock_create_truss_service,
+    mock_truss_handle,
+):
+    remote.push(
+        mock_truss_handle, "model_name", mock_truss_handle.truss_dir, publish=True
+    )
+
+    mock_create_truss_service.assert_called_once()
+    _, kwargs = mock_create_truss_service.call_args
+    assert kwargs.get("deploy_timeout_minutes") is None
+
+
+def test_push_integration_deploy_timeout_minutes_propagated(
+    custom_model_truss_dir_with_pre_and_post,
+    remote,
+    mock_baseten_requests,
+    mock_upload_truss,
+    mock_create_truss_service,
+    mock_truss_handle,
+):
+    remote.push(
+        mock_truss_handle,
+        "model_name",
+        mock_truss_handle.truss_dir,
+        publish=True,
+        environment="staging",
+        deploy_timeout_minutes=750,
+    )
+
+    mock_create_truss_service.assert_called_once()
+    _, kwargs = mock_create_truss_service.call_args
+    assert kwargs["deploy_timeout_minutes"] == 750
+    assert kwargs["environment"] == "staging"
+
+
+def test_api_push_integration_deploy_timeout_minutes_propagated(
+    custom_model_truss_dir_with_pre_and_post,
+    mock_remote_factory,
+    temp_trussrc_dir,
+    mock_available_config_names,
+    mock_truss_handle,
+):
+    from truss.api import push
+
+    push(
+        str(mock_truss_handle.truss_dir),
+        remote="baseten",
+        model_name="test_model",
+        deploy_timeout_minutes=1200,
+    )
+
+    # Verify the remote.push was called with deploy_timeout_minutes
+    mock_remote_factory.push.assert_called_once()
+    _, push_kwargs = mock_remote_factory.push.call_args
+    assert push_kwargs.get("deploy_timeout_minutes") == 1200

truss/tests/test_config.py

@@ -23,6 +23,7 @@ from truss.base.truss_config import (
     HTTPOptions,
     ModelCache,
     ModelRepo,
+    ModelRepoCacheInternal,
     Resources,
     Runtime,
     TransportKind,
@@ -292,7 +293,10 @@ def test_cache_internal_with_models(default_config):
     config = TrussConfig(
         python_version="py39",
         cache_internal=CacheInternal(
-            [ModelRepo(repo_id="test/model"), ModelRepo(repo_id="test/model2")]
+            [
+                ModelRepoCacheInternal(repo_id="test/model"),
+                ModelRepoCacheInternal(repo_id="test/model2"),
+            ]
         ),
     )
     new_config = default_config
@@ -305,11 +309,12 @@ def test_cache_internal_with_models(default_config):
 
 def test_huggingface_cache_single_model_default_revision(default_config):
     config = TrussConfig(
-        python_version="py39", model_cache=ModelCache([ModelRepo(repo_id="test/model")])
+        python_version="py39",
+        model_cache=ModelCache([ModelRepo(repo_id="test/model", use_volume=False)]),
     )
 
     new_config = default_config
-    new_config["model_cache"] = [{"repo_id": "test/model"}]
+    new_config["model_cache"] = [{"repo_id": "test/model", "use_volume": False}]
 
     assert new_config == config.to_dict(verbose=False)
     assert config.to_dict(verbose=True)["model_cache"][0].get("revision") is None
@@ -319,7 +324,9 @@ def test_huggingface_cache_single_model_non_default_revision_v1():
     config = TrussConfig(
         python_version="py39",
         requirements=[],
-        model_cache=ModelCache([ModelRepo(repo_id="test/model", revision="not-main")]),
+        model_cache=ModelCache(
+            [ModelRepo(repo_id="test/model", revision="not-main", use_volume=False)]
+        ),
     )
 
     assert config.to_dict(verbose=False)["model_cache"][0].get("revision") == "not-main"
@@ -330,16 +337,16 @@ def test_huggingface_cache_multiple_models_default_revision(default_config):
         python_version="py39",
         model_cache=ModelCache(
             [
-                ModelRepo(repo_id="test/model1", revision="main"),
-                ModelRepo(repo_id="test/model2"),
+                ModelRepo(repo_id="test/model1", revision="main", use_volume=False),
+                ModelRepo(repo_id="test/model2", use_volume=False),
             ]
         ),
     )
 
     new_config = default_config
     new_config["model_cache"] = [
-        {"repo_id": "test/model1", "revision": "main"},
-        {"repo_id": "test/model2"},
+        {"repo_id": "test/model1", "revision": "main", "use_volume": False},
+        {"repo_id": "test/model2", "use_volume": False},
     ]
 
     assert new_config == config.to_dict(verbose=False)
@@ -355,16 +362,18 @@ def test_huggingface_cache_multiple_models_mixed_revision(default_config):
         python_version="py39",
         model_cache=ModelCache(
             [
-                ModelRepo(repo_id="test/model1"),
-                ModelRepo(repo_id="test/model2", revision="not-main2"),
+                ModelRepo(repo_id="test/model1", use_volume=False),
+                ModelRepo(
+                    repo_id="test/model2", revision="not-main2", use_volume=False
+                ),
             ]
         ),
     )
 
     new_config = default_config
     new_config["model_cache"] = [
-        {"repo_id": "test/model1"},
-        {"repo_id": "test/model2", "revision": "not-main2"},
+        {"repo_id": "test/model1", "use_volume": False},
+        {"repo_id": "test/model2", "revision": "not-main2", "use_volume": False},
     ]
 
     assert new_config == config.to_dict(verbose=False)