truss 0.11.18rc500__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

@@ -186,6 +186,7 @@ def test_push_chain_atomic_with_chain_upload(
     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,
@@ -194,13 +195,18 @@ def test_push_chain_atomic_with_chain_upload(
         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")
@@ -239,6 +245,9 @@ def test_push_chain_atomic_without_chain_upload(
             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")

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