truss 0.11.6rc102__py3-none-any.whl → 0.11.24rc2__py3-none-any.whl

truss/tests/remote/baseten/test_api.py

@@ -165,9 +165,9 @@ def test_create_model_version_from_truss(mock_post, baseten_api):
         "config_str",
         "semver_bump",
         b10_types.TrussUserEnv.collect(),
-        False,
-        "deployment_name",
-        "production",
+        preserve_previous_prod_deployment=False,
+        deployment_name="deployment_name",
+        environment="production",
     )
 
     gql_mutation = mock_post.call_args[1]["json"]["query"]
@@ -182,6 +182,7 @@ def test_create_model_version_from_truss(mock_post, baseten_api):
     assert 'name: "deployment_name"' in gql_mutation
     assert 'environment_name: "production"' in gql_mutation
     assert "preserve_env_instance_type: true" in gql_mutation
+    assert "deploy_timeout_minutes: " not in gql_mutation
 
 
 @mock.patch("requests.post", return_value=mock_create_model_version_response())
@@ -211,6 +212,7 @@ def test_create_model_version_from_truss_does_not_send_deployment_name_if_not_sp
     assert " name: " not in gql_mutation
     assert "environment_name: " not in gql_mutation
     assert "preserve_env_instance_type: false" in gql_mutation
+    assert "deploy_timeout_minutes: " not in gql_mutation
 
 
 @mock.patch("requests.post", return_value=mock_create_model_version_response())
@@ -242,6 +244,57 @@ def test_create_model_version_from_truss_does_not_scale_old_prod_to_zero_if_keep
     assert " name: " not in gql_mutation
     assert 'environment_name: "staging"' in gql_mutation
     assert "preserve_env_instance_type: true" in gql_mutation
+    assert "deploy_timeout_minutes: " not in gql_mutation
+
+
+@mock.patch("requests.post", return_value=mock_create_model_version_response())
+def test_create_model_version_from_truss_with_deploy_timeout_minutes(
+    mock_post, baseten_api
+):
+    baseten_api.create_model_version_from_truss(
+        "model_id",
+        "s3key",
+        "config_str",
+        "semver_bump",
+        b10_types.TrussUserEnv.collect(),
+        preserve_previous_prod_deployment=False,
+        deployment_name="deployment_name",
+        environment="production",
+        deploy_timeout_minutes=300,
+    )
+
+    gql_mutation = mock_post.call_args[1]["json"]["query"]
+    assert 'model_id: "model_id"' in gql_mutation
+    assert 's3_key: "s3key"' in gql_mutation
+    assert 'config: "config_str"' in gql_mutation
+    assert 'semver_bump: "semver_bump"' in gql_mutation
+    assert {
+        "trussUserEnv": b10_types.TrussUserEnv.collect().model_dump_json()
+    } == mock_post.call_args[1]["json"]["variables"]
+    assert "scale_down_old_production: true" in gql_mutation
+    assert 'name: "deployment_name"' in gql_mutation
+    assert 'environment_name: "production"' in gql_mutation
+    assert "preserve_env_instance_type: true" in gql_mutation
+    assert "deploy_timeout_minutes: 300" in gql_mutation
+
+
+@mock.patch("requests.post", return_value=mock_create_model_version_response())
+def test_create_model_version_from_truss_with_deploy_timeout_minutes_zero(
+    mock_post, baseten_api
+):
+    """Test that deploy_timeout_minutes of 0 is handled correctly"""
+    baseten_api.create_model_version_from_truss(
+        "model_id",
+        "s3key",
+        "config_str",
+        "semver_bump",
+        b10_types.TrussUserEnv.collect(),
+        preserve_previous_prod_deployment=False,
+        deploy_timeout_minutes=0,
+    )
+
+    gql_mutation = mock_post.call_args[1]["json"]["query"]
+    assert "deploy_timeout_minutes: 0" in gql_mutation
 
 
 @mock.patch("requests.post", return_value=mock_create_model_response())
@@ -332,6 +385,48 @@ def test_create_development_model_from_truss_with_allow_truss_download(
         "trussUserEnv": b10_types.TrussUserEnv.collect().model_dump_json()
     } == mock_post.call_args[1]["json"]["variables"]
     assert "allow_truss_download: false" in gql_mutation
+    assert "deploy_timeout_minutes: " not in gql_mutation
+
+
+@mock.patch("requests.post", return_value=mock_create_development_model_response())
+def test_create_development_model_from_truss_with_deploy_timeout_minutes(
+    mock_post, baseten_api
+):
+    baseten_api.create_development_model_from_truss(
+        "model_name",
+        "s3key",
+        "config_str",
+        b10_types.TrussUserEnv.collect(),
+        allow_truss_download=False,
+        deploy_timeout_minutes=300,
+    )
+
+    gql_mutation = mock_post.call_args[1]["json"]["query"]
+    assert 'name: "model_name"' in gql_mutation
+    assert 's3_key: "s3key"' in gql_mutation
+    assert 'config: "config_str"' in gql_mutation
+    assert {
+        "trussUserEnv": b10_types.TrussUserEnv.collect().model_dump_json()
+    } == mock_post.call_args[1]["json"]["variables"]
+    assert "allow_truss_download: false" in gql_mutation
+    assert "deploy_timeout_minutes: 300" in gql_mutation
+
+
+@mock.patch("requests.post", return_value=mock_create_development_model_response())
+def test_create_development_model_from_truss_with_deploy_timeout_minutes_zero(
+    mock_post, baseten_api
+):
+    """Test that deploy_timeout_minutes of 0 is handled correctly"""
+    baseten_api.create_development_model_from_truss(
+        "model_name",
+        "s3key",
+        "config_str",
+        b10_types.TrussUserEnv.collect(),
+        deploy_timeout_minutes=0,
+    )
+
+    gql_mutation = mock_post.call_args[1]["json"]["query"]
+    assert "deploy_timeout_minutes: 0" in gql_mutation
 
 
 @mock.patch("requests.post", return_value=mock_deploy_chain_deployment_response())
@@ -357,6 +452,30 @@ def test_deploy_chain_deployment(mock_post, baseten_api):
     assert 'chain_id: "chain_id"' in gql_mutation
     assert "dependencies:" in gql_mutation
     assert "entrypoint:" in gql_mutation
+    assert "deployment_name" not in gql_mutation
+
+
+@mock.patch("requests.post", return_value=mock_deploy_chain_deployment_response())
+def test_deploy_chain_deployment_with_deployment_name(mock_post, baseten_api):
+    baseten_api.deploy_chain_atomic(
+        environment="production",
+        chain_id="chain_id",
+        dependencies=[],
+        entrypoint=ChainletDataAtomic(
+            name="chainlet-1",
+            oracle=OracleData(
+                model_name="model-1",
+                s3_key="s3-key-1",
+                encoded_config_str="encoded-config-str-1",
+            ),
+        ),
+        truss_user_env=b10_types.TrussUserEnv.collect(),
+        deployment_name="chain-deployment-name",
+    )
+
+    gql_mutation = mock_post.call_args[1]["json"]["query"]
+
+    assert 'deployment_name: "chain-deployment-name"' in gql_mutation
 
 
 @mock.patch("requests.post", return_value=mock_deploy_chain_deployment_response())

truss/tests/remote/baseten/test_chain_upload.py

@@ -0,0 +1,294 @@
+import pathlib
+import tempfile
+from unittest.mock import Mock, patch
+
+import pytest
+
+from truss.remote.baseten import custom_types as b10_types
+from truss.remote.baseten.api import BasetenApi
+from truss.remote.baseten.core import upload_chain_artifact
+from truss.remote.baseten.remote import BasetenRemote
+
+
+@pytest.fixture
+def mock_push_data():
+    """Fixture providing mock push data for tests."""
+    mock_push_data = Mock()
+    mock_push_data.model_name = "test-model"
+    mock_push_data.s3_key = "models/test-key"
+    mock_push_data.encoded_config_str = "encoded_config"
+    mock_push_data.is_draft = False
+    mock_push_data.model_id = "model-id"
+    mock_push_data.version_name = None
+    return mock_push_data
+
+
+@pytest.fixture
+def mock_remote_context():
+    """Fixture providing mock remote and context managers for tests."""
+    api = Mock(spec=BasetenApi)
+
+    remote = BasetenRemote("https://test.baseten.co", "test-api-key")
+    remote._api = api
+
+    chain_name = "test-chain"
+    entrypoint_artifact = Mock()
+    entrypoint_artifact.truss_dir = "/path/to/truss"
+    entrypoint_artifact.display_name = "entrypoint"
+
+    dependency_artifacts = []
+    truss_user_env = Mock()
+    chain_root = pathlib.Path("/path/to/chain")
+
+    with patch.object(remote, "_prepare_push") as mock_prepare_push:
+        with patch("truss.remote.baseten.remote.truss_build.load") as mock_load:
+            mock_truss_handle = Mock()
+            mock_truss_handle.spec.config.model_name = "test-model"
+            mock_load.return_value = mock_truss_handle
+
+            yield {
+                "remote": remote,
+                "api": api,
+                "chain_name": chain_name,
+                "entrypoint_artifact": entrypoint_artifact,
+                "dependency_artifacts": dependency_artifacts,
+                "truss_user_env": truss_user_env,
+                "chain_root": chain_root,
+                "mock_prepare_push": mock_prepare_push,
+                "mock_load": mock_load,
+                "mock_truss_handle": mock_truss_handle,
+            }
+
+
+def test_get_blob_credentials_for_chain():
+    """Test that get_blob_credentials works correctly for chain blob type using GraphQL."""
+    mock_graphql_response = {
+        "data": {
+            "chain_s3_upload_credentials": {
+                "s3_bucket": "test-chain-bucket",
+                "s3_key": "chains/test-uuid/chain.tgz",
+                "aws_access_key_id": "test_access_key",
+                "aws_secret_access_key": "test_secret_key",
+                "aws_session_token": "test_session_token",
+            }
+        }
+    }
+
+    # Create a real API instance and mock the GraphQL call
+    mock_auth_service = Mock()
+    mock_auth_service.authenticate.return_value = Mock(value="test-token")
+    api = BasetenApi("https://test.baseten.co", mock_auth_service)
+    with patch.object(api, "_post_graphql_query") as mock_graphql:
+        mock_graphql.return_value = mock_graphql_response
+
+        result = api.get_chain_s3_upload_credentials()
+
+        assert result.s3_bucket == "test-chain-bucket"
+        assert result.s3_key == "chains/test-uuid/chain.tgz"
+        assert result.aws_access_key_id == "test_access_key"
+        assert result.aws_secret_access_key == "test_secret_key"
+        assert result.aws_session_token == "test_session_token"
+
+        mock_graphql.assert_called_once()
+        call_args = mock_graphql.call_args
+        assert "chain_s3_upload_credentials" in call_args[0][0]
+
+
+def test_get_blob_credentials_for_other_types_uses_rest():
+    """Test that get_blob_credentials uses REST API for non-chain blob types."""
+    mock_response = {
+        "s3_bucket": "test-bucket",
+        "s3_key": "test-key",
+        "creds": {
+            "aws_access_key_id": "test_access_key",
+            "aws_secret_access_key": "test_secret_key",
+            "aws_session_token": "test_session_token",
+        },
+    }
+
+    mock_auth_service = Mock()
+    mock_auth_service.authenticate.return_value = Mock(value="test-token")
+    api = BasetenApi("https://test.baseten.co", mock_auth_service)
+    with patch.object(api, "_rest_api_client") as mock_client, patch.object(
+        api, "_post_graphql_query"
+    ) as mock_graphql:
+        mock_client.get.return_value = mock_response
+
+        result = api.get_blob_credentials(b10_types.BlobType.MODEL)
+
+        assert result["s3_bucket"] == "test-bucket"
+        assert result["s3_key"] == "test-key"
+
+        mock_client.get.assert_called_once_with("v1/blobs/credentials/model")
+        mock_graphql.assert_not_called()
+
+
+@patch("truss.remote.baseten.core.multipart_upload_boto3")
+def test_upload_chain_artifact_function(mock_multipart_upload):
+    """Test the upload_chain_artifact function."""
+    # Mock ChainUploadCredentials object
+    mock_credentials = Mock()
+    mock_credentials.s3_bucket = "test-chain-bucket"
+    mock_credentials.s3_key = "chains/test-uuid/chain.tgz"
+    mock_credentials.aws_credentials = Mock()
+    mock_credentials.aws_credentials.model_dump.return_value = {
+        "aws_access_key_id": "test_access_key",
+        "aws_secret_access_key": "test_secret_key",
+        "aws_session_token": "test_session_token",
+    }
+
+    api = Mock(spec=BasetenApi)
+    api.get_chain_s3_upload_credentials.return_value = mock_credentials
+
+    with tempfile.NamedTemporaryFile(suffix=".tgz", delete=False) as temp_file:
+        temp_file.write(b"test chain content")
+        temp_file.flush()
+
+        result = upload_chain_artifact(api, temp_file, None)
+
+        assert result == "chains/test-uuid/chain.tgz"
+
+        api.get_chain_s3_upload_credentials.assert_called_once_with()
+
+        mock_multipart_upload.assert_called_once()
+        call_args = mock_multipart_upload.call_args
+        assert call_args[0][0] == temp_file.name  # file path
+        assert call_args[0][1] == "test-chain-bucket"  # bucket
+        assert call_args[0][2] == "chains/test-uuid/chain.tgz"  # key
+        assert call_args[0][3] == {  # credentials
+            "aws_access_key_id": "test_access_key",
+            "aws_secret_access_key": "test_secret_key",
+            "aws_session_token": "test_session_token",
+        }
+
+
+@patch("truss.remote.baseten.remote.upload_chain_artifact")
+@patch("truss.remote.baseten.remote.archive_dir")
+@patch("truss.remote.baseten.remote.create_chain_atomic")
+def test_push_chain_atomic_with_chain_upload(
+    mock_create_chain_atomic,
+    mock_archive_dir,
+    mock_upload_chain_artifact,
+    mock_push_data,
+    mock_remote_context,
+):
+    """Test that push_chain_atomic uploads raw chain artifact when chain_root is provided."""
+    mock_create_chain_atomic.return_value = Mock()
+    mock_archive_dir.return_value = Mock()
+    mock_upload_chain_artifact.return_value = "chains/test-uuid/chain.tgz"
+
+    context = mock_remote_context
+    remote = context["remote"]
+    chain_name = context["chain_name"]
+    entrypoint_artifact = context["entrypoint_artifact"]
+    dependency_artifacts = context["dependency_artifacts"]
+    truss_user_env = context["truss_user_env"]
+    chain_root = context["chain_root"]
+
+    context["mock_prepare_push"].return_value = mock_push_data
+    deployment_name = "custom_deployment"
+
+    result = remote.push_chain_atomic(
+        chain_name=chain_name,
+        entrypoint_artifact=entrypoint_artifact,
+        dependency_artifacts=dependency_artifacts,
+        truss_user_env=truss_user_env,
+        chain_root=chain_root,
+        publish=True,
+        deployment_name=deployment_name,
+    )
+    assert result == mock_create_chain_atomic.return_value
+
+    mock_archive_dir.assert_called_once_with(dir=chain_root, progress_bar=None)
+    mock_upload_chain_artifact.assert_called_once()
+    mock_create_chain_atomic.assert_called_once()
+    create_kwargs = mock_create_chain_atomic.call_args.kwargs
+    assert create_kwargs["deployment_name"] == deployment_name
+
+    prepare_kwargs = context["mock_prepare_push"].call_args.kwargs
+    assert prepare_kwargs["deployment_name"] == deployment_name
+
+
+@patch("truss.remote.baseten.remote.create_chain_atomic")
+def test_push_chain_atomic_without_chain_upload(
+    mock_create_chain_atomic, mock_push_data, mock_remote_context
+):
+    """Test that push_chain_atomic skips chain upload when chain_root is None."""
+    mock_create_chain_atomic.return_value = Mock()
+
+    context = mock_remote_context
+    remote = context["remote"]
+    chain_name = context["chain_name"]
+    entrypoint_artifact = context["entrypoint_artifact"]
+    dependency_artifacts = context["dependency_artifacts"]
+    truss_user_env = context["truss_user_env"]
+
+    context["mock_prepare_push"].return_value = mock_push_data
+
+    with patch("truss.remote.baseten.remote.upload_chain_artifact") as mock_upload:
+        with patch(
+            "truss.remote.baseten.core.create_tar_with_progress_bar"
+        ) as mock_tar:
+            # Call push_chain_atomic without chain_root
+            result = remote.push_chain_atomic(
+                chain_name=chain_name,
+                entrypoint_artifact=entrypoint_artifact,
+                dependency_artifacts=dependency_artifacts,
+                truss_user_env=truss_user_env,
+                chain_root=None,  # No chain root
+                publish=True,
+            )
+
+            assert result
+            # Verify chain artifact upload was NOT called
+            mock_tar.assert_not_called()
+            mock_upload.assert_not_called()
+
+            mock_create_chain_atomic.assert_called_once()
+            create_kwargs = mock_create_chain_atomic.call_args.kwargs
+            assert "deployment_name" in create_kwargs
+            assert create_kwargs["deployment_name"] is None
+
+
+@patch("truss.remote.baseten.core.multipart_upload_boto3")
+def test_upload_chain_artifact_error_handling(mock_multipart_upload):
+    """Test error handling in upload_chain_artifact."""
+    # Mock API to raise an exception
+    api = Mock(spec=BasetenApi)
+    api.get_chain_s3_upload_credentials.side_effect = Exception("API Error")
+
+    with tempfile.NamedTemporaryFile(suffix=".tgz") as temp_file:
+        # Should raise the exception
+        with pytest.raises(Exception, match="API Error"):
+            upload_chain_artifact(api, temp_file, None)
+
+
+def test_upload_chain_artifact_credentials_extraction():
+    """Test that credentials are properly extracted from API response."""
+    # Mock ChainUploadCredentials object
+    mock_credentials = Mock()
+    mock_credentials.s3_bucket = "test-bucket"
+    mock_credentials.s3_key = "chains/test-uuid/chain.tgz"
+    mock_credentials.aws_credentials = Mock()
+    mock_credentials.aws_credentials.model_dump.return_value = {
+        "aws_access_key_id": "access_key",
+        "aws_secret_access_key": "secret_key",
+        "aws_session_token": "session_token",
+    }
+
+    api = Mock(spec=BasetenApi)
+    api.get_chain_s3_upload_credentials.return_value = mock_credentials
+
+    with patch("truss.remote.baseten.core.multipart_upload_boto3") as mock_upload:
+        with tempfile.NamedTemporaryFile(suffix=".tgz") as temp_file:
+            upload_chain_artifact(api, temp_file, None)
+
+            call_args = mock_upload.call_args
+            credentials = call_args[0][3]
+
+            assert credentials == {
+                "aws_access_key_id": "access_key",
+                "aws_secret_access_key": "secret_key",
+                "aws_session_token": "session_token",
+            }
+            assert "extra_field" not in credentials

truss/tests/remote/baseten/test_core.py

@@ -187,6 +187,7 @@ def test_create_truss_service_handles_existing_model(inputs):
     _, kwargs = api.create_model_version_from_truss.call_args
     for k, v in inputs.items():
         assert kwargs[k] == v
+    assert kwargs.get("deploy_timeout_minutes") is None
 
 
 @pytest.mark.parametrize("allow_truss_download", [True, False])
@@ -761,3 +762,88 @@ def test_get_training_job_logs_with_pagination_default_batch_size(baseten_api):
     query_params = call_args[0][2]  # query_params
 
     assert query_params["limit"] == MAX_BATCH_SIZE
+
+
+def test_create_truss_service_passes_deploy_timeout_minutes():
+    """Test that deploy_timeout_minutes is passed through to create_model_version_from_truss"""
+    api = MagicMock()
+    return_value = {
+        "id": "model_version_id",
+        "oracle": {"id": "model_id", "hostname": "hostname"},
+        "instance_type": {"name": "1x2"},
+    }
+    api.create_model_version_from_truss.return_value = return_value
+    version_handle = create_truss_service(
+        api,
+        "model_name",
+        "s3_key",
+        "config",
+        b10_types.TrussUserEnv.collect(),
+        is_draft=False,
+        model_id="model_id",
+        environment="staging",
+        deploy_timeout_minutes=600,
+    )
+
+    assert version_handle.version_id == "model_version_id"
+    assert version_handle.model_id == "model_id"
+    api.create_model_version_from_truss.assert_called_once()
+    _, kwargs = api.create_model_version_from_truss.call_args
+    assert kwargs["deploy_timeout_minutes"] == 600
+
+
+def test_create_truss_service_passes_deploy_timeout_minutes_with_other_params():
+    """Test that deploy_timeout_minutes works correctly with other parameters like preserve_env_instance_type"""
+    api = MagicMock()
+    return_value = {
+        "id": "model_version_id",
+        "oracle": {"id": "model_id", "hostname": "hostname"},
+        "instance_type": {"name": "1x2"},
+    }
+    api.create_model_version_from_truss.return_value = return_value
+    version_handle = create_truss_service(
+        api,
+        "model_name",
+        "s3_key",
+        "config",
+        b10_types.TrussUserEnv.collect(),
+        is_draft=False,
+        model_id="model_id",
+        environment="production",
+        preserve_env_instance_type=False,
+        deploy_timeout_minutes=900,
+    )
+
+    assert version_handle.version_id == "model_version_id"
+    api.create_model_version_from_truss.assert_called_once()
+    _, kwargs = api.create_model_version_from_truss.call_args
+    assert kwargs["deploy_timeout_minutes"] == 900
+    assert kwargs["preserve_env_instance_type"] is False
+    assert kwargs["environment"] == "production"
+
+
+def test_create_truss_service_passes_deploy_timeout_minutes_for_development_model():
+    """Test that deploy_timeout_minutes is passed through to create_development_model_from_truss"""
+    api = MagicMock()
+    return_value = {
+        "id": "model_version_id",
+        "oracle": {"id": "model_id", "hostname": "hostname"},
+        "instance_type": {"name": "1x2"},
+    }
+    api.create_development_model_from_truss.return_value = return_value
+    version_handle = create_truss_service(
+        api,
+        "model_name",
+        "s3_key",
+        "config",
+        b10_types.TrussUserEnv.collect(),
+        is_draft=True,
+        model_id=None,
+        deploy_timeout_minutes=600,
+    )
+
+    assert version_handle.version_id == "model_version_id"
+    assert version_handle.model_id == "model_id"
+    api.create_development_model_from_truss.assert_called_once()
+    _, kwargs = api.create_development_model_from_truss.call_args
+    assert kwargs["deploy_timeout_minutes"] == 600