fix: Python 3.14 tuner tags (#6554) and tuner source staging bucket (#6555)

sagemaker-core/src/sagemaker/core/tools/constants.py

@@ -53,6 +53,16 @@
 BASIC_RETURN_TYPES = {"str", "int", "bool", "float", "datetime.datetime"}
 
 SHAPE_DAG_FILE_PATH = os.getcwd() + "/src/sagemaker/core/utils/code_injection/shape_dag.py"
+
+# Shape names that a resource class of the same name shadows inside resources.py.
+# resources.py star-imports the shapes, then later defines ``class Tag(Base)`` (the
+# ListTags resource). An unqualified ``Tag`` annotation written before that class means
+# the shape only while annotations are evaluated eagerly; Python 3.14 evaluates them
+# lazily (PEP 649), so it resolves to the resource class and rejects shape instances
+# (issue #6554). Resource codegen emits these names as ``shapes.<Name>``. Other shared
+# names (Device, Workteam, ...) are only referenced after their resource class, so they
+# already resolve to the resource class on every Python version and are left alone.
+RESOURCE_SHADOWED_SHAPE_NAMES = frozenset({"Tag"})
 PYTHON_TYPES_TO_BASIC_JSON_TYPES = {
     "str": "string",
     "StrPipeVar": "string",

sagemaker-core/src/sagemaker/core/tools/resources_codegen.py

@@ -29,6 +29,7 @@
     LOGGER_STRING,
     CONFIG_SCHEMA_FILE_NAME,
     PYTHON_TYPES_TO_BASIC_JSON_TYPES,
+    RESOURCE_SHADOWED_SHAPE_NAMES,
     CONFIGURABLE_ATTRIBUTE_SUBSTRINGS,
     RESOURCE_WITH_LOGS,
     DEFAULT_TIMEOUT_MESSAGE,
@@ -154,7 +155,7 @@ def __init__(self, service_json: dict):
 
         # Initialize the resources and shapes extractors
         self.resources_extractor = ResourcesExtractor()
-        self.shapes_extractor = ShapesExtractor()
+        self.shapes_extractor = ShapesExtractor(qualified_shape_names=RESOURCE_SHADOWED_SHAPE_NAMES)
 
         # Extract the resources plan and shapes DAG
         self.resources_plan = self.resources_extractor.get_resource_plan()
@@ -196,6 +197,9 @@ def generate_imports(self) -> str:
             "from rich.status import Status",
             "from rich.style import Style",
             "from sagemaker.core.shapes import *",
+            # Qualified access for shape names a resource class shadows (see
+            # RESOURCE_SHADOWED_SHAPE_NAMES), e.g. ``shapes.Tag`` vs resource ``Tag``.
+            "from sagemaker.core import shapes",
             "from sagemaker.core.helper.pipeline_variable import StrPipeVar",
             "from sagemaker.core.utils.code_injection.codec import transform",
             "from sagemaker.core.utils.code_injection.constants import Color",

sagemaker-core/src/sagemaker/core/tools/shapes_extractor.py

@@ -34,12 +34,20 @@
 class ShapesExtractor:
     """Extracts the shapes to DAG structure."""
 
-    def __init__(self, combined_shapes: Optional[dict] = None):
+    def __init__(
+        self,
+        combined_shapes: Optional[dict] = None,
+        qualified_shape_names: frozenset = frozenset(),
+    ):
         """Initializes a new instance of the ShapesExtractor class.
 
         :param combined_shapes: All the shapes put together from all Sagemaker Service JSONs
+        :param qualified_shape_names: Structure shape names to emit as ``shapes.<Name>``
+            in type annotations. Used by resource codegen for shapes that a resource
+            class of the same name shadows in resources.py. Empty for shape codegen.
         """
         self.combined_shapes = combined_shapes or load_combined_shapes_data()
+        self.qualified_shape_names = qualified_shape_names
 
         self.shape_dag = self.get_shapes_dag()
         with open(SHAPE_DAG_FILE_PATH, "w") as f:
@@ -121,6 +129,12 @@ def get_shapes_dag(self):
                 _dag[shape]["value_type"] = _all_shapes[_map_value_shape]["type"]
         return _dag
 
+    def _structure_type(self, shape_name):
+        """Return the annotation text for a structure shape, qualified when shadowed."""
+        if shape_name in self.qualified_shape_names:
+            return f"shapes.{shape_name}"
+        return shape_name
+
     def _evaluate_list_type(self, member_shape):
         list_shape_name = member_shape["member"]["shape"]
         list_shape_member = self.combined_shapes[list_shape_name]
@@ -134,7 +148,7 @@ def _evaluate_list_type(self, member_shape):
             if list_shape_name == "SearchExpression":
                 member_type = f"List['{list_shape_name}']"
             else:
-                member_type = f"List[{list_shape_name}]"
+                member_type = f"List[{self._structure_type(list_shape_name)}]"
         elif list_shape_type in BASIC_JSON_TYPES_TO_PYTHON_TYPES.keys():
             member_type = f"List[{BASIC_JSON_TYPES_TO_PYTHON_TYPES[list_shape_type]}]"
         else:
@@ -158,7 +172,7 @@ def _evaluate_map_type(self, member_shape):
         if map_value_shape_type == "structure":
             member_type = (
                 f"Dict[{BASIC_JSON_TYPES_TO_PYTHON_TYPES[map_key_shape_type]}, "
-                f"{map_value_shape_name}]"
+                f"{self._structure_type(map_value_shape_name)}]"
             )
         elif map_value_shape_type == "list":
             member_type = (
@@ -238,7 +252,7 @@ def generate_shape_members(self, shape, required_override=()):
                 member_shape = self.combined_shapes[member_shape_name]
                 member_shape_type = member_shape["type"]
                 if member_shape_type == "structure":
-                    member_type = member_shape_name
+                    member_type = self._structure_type(member_shape_name)
                 elif member_shape_type == "list":
                     member_type = self._evaluate_list_type(member_shape)
                 elif member_shape_type == "map":

sagemaker-core/tests/unit/generated/test_resources_shadowed_shapes.py

@@ -0,0 +1,102 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Guard resource annotations against shape names shadowed by resource classes.
+
+resources.py star-imports every shape and then defines resource classes, some of which
+share a shape's name (``Tag``). Before Python 3.14 an annotation is evaluated when its
+function or class is defined, so a bare ``Tag`` written above ``class Tag`` means the
+shape. Python 3.14 evaluates annotations lazily (PEP 649), so the same text resolves to
+the resource class and rejects shape instances (issue #6554).
+"""
+
+import ast
+import inspect
+import typing
+
+import pytest
+from pydantic import TypeAdapter
+
+from sagemaker.core import resources, shapes
+
+
+def _resource_methods_with_tags():
+    for _, cls in inspect.getmembers(resources, inspect.isclass):
+        if cls.__module__ != resources.__name__:
+            continue
+        for method_name in ("create", "add_tags"):
+            method = cls.__dict__.get(method_name)
+            if method is None:
+                continue
+            func = method.__func__ if isinstance(method, classmethod) else method
+            func = inspect.unwrap(func)
+            if "tags" in inspect.signature(func).parameters:
+                yield f"{cls.__name__}.{method_name}", func
+
+
+TAGGED_METHODS = sorted(_resource_methods_with_tags())
+
+
+def test_no_shape_name_used_before_its_shadowing_resource_class():
+    """A bare annotation name must not change meaning between eager and lazy evaluation."""
+    tree = ast.parse(inspect.getsource(resources))
+    shape_names = {name for name, obj in vars(shapes).items() if inspect.isclass(obj)}
+    shadowing_class_line = {
+        node.name: node.lineno
+        for node in tree.body
+        if isinstance(node, ast.ClassDef) and node.name in shape_names and node.name != "Base"
+    }
+
+    def annotations(node):
+        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+            args = node.args
+            for arg in args.posonlyargs + args.args + args.kwonlyargs:
+                yield arg.annotation
+            yield node.returns
+        elif isinstance(node, ast.AnnAssign):
+            yield node.annotation
+
+    offenders = set()
+    for node in ast.walk(tree):
+        for annotation in annotations(node):
+            if annotation is None:
+                continue
+            for name in ast.walk(annotation):
+                if (
+                    isinstance(name, ast.Name)
+                    and name.id in shadowing_class_line
+                    and name.lineno < shadowing_class_line[name.id]
+                ):
+                    offenders.add((name.id, name.lineno))
+
+    assert not offenders, (
+        "Annotations reference a shape name that a later resource class shadows; add the "
+        f"name to RESOURCE_SHADOWED_SHAPE_NAMES and regenerate: {sorted(offenders)[:10]}"
+    )
+
+
+def test_tagged_methods_found():
+    assert len(TAGGED_METHODS) > 50
+    assert "HyperParameterTuningJob.create" in dict(TAGGED_METHODS)
+
+
+@pytest.mark.parametrize("name,func", TAGGED_METHODS, ids=[name for name, _ in TAGGED_METHODS])
+def test_tags_annotation_resolves_to_shape_tag(name, func):
+    """Every generated ``tags`` parameter takes the shape Tag on every Python version."""
+    tags_type = typing.get_type_hints(func)["tags"]
+    adapter = TypeAdapter(tags_type)
+
+    validated = adapter.validate_python([shapes.Tag(key="project", value="p1")])
+    assert type(validated[0]) is shapes.Tag
+
+    validated = adapter.validate_python([{"key": "project", "value": "p1"}])
+    assert type(validated[0]) is shapes.Tag

sagemaker-core/tests/unit/tools/test_shapes_extractor.py

@@ -492,3 +492,48 @@ def test_other_members_stay_required(self, extractor):
         assert members["capture_status"] == "StrPipeVar"
         assert members["current_sampling_percentage"] == "int"
         assert members["destination_s3_uri"] == "StrPipeVar"
+
+
+class TestQualifiedShapeNames:
+    """Shapes that a resource class shadows are emitted as ``shapes.<Name>`` (issue #6554)."""
+
+    _SHAPES = {
+        "Holder": {
+            "type": "structure",
+            "members": {
+                "Tag": {"shape": "Tag"},
+                "Tags": {"shape": "TagList"},
+                "TagMap": {"shape": "TagMap"},
+                "Other": {"shape": "Other"},
+            },
+        },
+        "Tag": {"type": "structure", "members": {}},
+        "Other": {"type": "structure", "members": {}},
+        "TagList": {"type": "list", "member": {"shape": "Tag"}},
+        "TagMap": {"type": "map", "key": {"shape": "String"}, "value": {"shape": "Tag"}},
+        "String": {"type": "string"},
+    }
+
+    @patch("sagemaker.core.tools.shapes_extractor.reformat_file_with_black")
+    def test_qualifies_only_listed_shapes(self, mock_reformat):
+        with patch("builtins.open", create=True):
+            extractor = ShapesExtractor(
+                combined_shapes=self._SHAPES, qualified_shape_names=frozenset({"Tag"})
+            )
+
+        members = extractor.generate_shape_members("Holder")
+
+        assert members["tag"] == "Optional[shapes.Tag] = Unassigned()"
+        assert members["tags"] == "Optional[List[shapes.Tag]] = Unassigned()"
+        assert members["tag_map"] == "Optional[Dict[StrPipeVar, shapes.Tag]] = Unassigned()"
+        assert members["other"] == "Optional[Other] = Unassigned()"
+
+    @patch("sagemaker.core.tools.shapes_extractor.reformat_file_with_black")
+    def test_default_leaves_names_unqualified(self, mock_reformat):
+        """Shape codegen (no qualified names) must keep emitting bare names."""
+        with patch("builtins.open", create=True):
+            extractor = ShapesExtractor(combined_shapes=self._SHAPES)
+
+        members = extractor.generate_shape_members("Holder")
+
+        assert members["tags"] == "Optional[List[Tag]] = Unassigned()"

sagemaker-train/src/sagemaker/train/tuner.py

@@ -577,8 +577,16 @@ def _build_driver_and_code_channels(cls, model_trainer):
             import tempfile
 
             session = model_trainer.sagemaker_session
-            bucket = session.default_bucket()
-            s3_key = f"{key_prefix}/sourcedir/sourcedir.tar.gz"
+            # Resolve bucket and prefix the same way as create_input_data_channel, so the
+            # source lands next to sm_drivers and honors default_bucket_prefix and the
+            # output-path fallback when the role can't write the default bucket (#6555).
+            source_key_prefix = f"{key_prefix}/sourcedir"
+            if session.default_bucket_prefix:
+                source_key_prefix = f"{session.default_bucket_prefix}/{source_key_prefix}"
+            bucket, staging_prefix = model_trainer._resolve_staging_bucket()
+            if staging_prefix:
+                source_key_prefix = f"{staging_prefix}/{source_key_prefix}"
+            s3_key = f"{source_key_prefix}/sourcedir.tar.gz"
 
             with tempfile.NamedTemporaryFile(suffix=".tar.gz", delete=False) as tmp:
                 tar_path = tmp.name

sagemaker-train/tests/unit/train/test_model_trainer.py

@@ -2795,10 +2795,9 @@ def test_boto3_request_uses_boto_style_tag_keys(model_trainer):
 def test_interactive_request_keeps_tags_training_job_create_accepts(model_trainer):
     """The interactive request carries dicts, not Tag models.
 
-    ``TrainingJob.create`` annotates ``tags`` with ``sagemaker.core.resources.Tag``,
-    which shadows the shape of the same name, so a ``sagemaker.core.shapes.Tag``
-    instance fails its validation. The dicts are revalidated into that model instead,
-    and come out boto-style on the wire.
+    ``TrainingJob.create`` validates ``tags`` as ``sagemaker.core.shapes.Tag`` (it used
+    to resolve to the shadowing ``sagemaker.core.resources.Tag``, see issue #6554). The
+    dicts are revalidated into that model and come out boto-style on the wire.
     """
     model_trainer.tags = [TAG]
 

sagemaker-train/tests/unit/train/test_tuner_driver_channels.py

@@ -89,6 +89,7 @@ def _mock_model_trainer(**overrides):
     trainer.tags = []
     trainer.source_code = None
     trainer.configure_mock(**{"networking": None})
+    trainer._resolve_staging_bucket.return_value = ("my-bucket", None)
     for k, v in overrides.items():
         setattr(trainer, k, v)
     return trainer
@@ -316,6 +317,33 @@ def test_uploads_sourcedir_without_kms_key(self, tmp_path):
 
         assert mock_s3_client.upload_file.call_args.kwargs["ExtraArgs"] is None
 
+    def test_uploads_sourcedir_to_resolved_staging_bucket(self, tmp_path):
+        """Source upload must follow the trainer's staging-bucket fallback (issue #6555)."""
+        trainer, mock_s3_client = self._upload_trainer(tmp_path, None)
+        trainer._resolve_staging_bucket.return_value = ("output-bucket", "jobs")
+
+        HyperparameterTuner._build_driver_and_code_channels(trainer)
+
+        trainer._resolve_staging_bucket.assert_called_once_with()
+        trainer.sagemaker_session.default_bucket.assert_not_called()
+        _, bucket, key = mock_s3_client.upload_file.call_args.args
+        assert bucket == "output-bucket"
+        assert key.startswith("jobs/test-tuning/")
+        assert key.endswith("/input/sourcedir/sourcedir.tar.gz")
+        assert trainer.hyperparameters["sagemaker_submit_directory"] == f"s3://output-bucket/{key}"
+
+    def test_uploads_sourcedir_with_default_bucket_prefix(self, tmp_path):
+        """Source upload must honor the session's default_bucket_prefix (issue #6555)."""
+        trainer, mock_s3_client = self._upload_trainer(tmp_path, None)
+        trainer.sagemaker_session.default_bucket_prefix = "team-prefix"
+
+        HyperparameterTuner._build_driver_and_code_channels(trainer)
+
+        _, bucket, key = mock_s3_client.upload_file.call_args.args
+        assert bucket == "my-bucket"
+        assert key.startswith("team-prefix/test-tuning/")
+        assert key.endswith("/sourcedir/sourcedir.tar.gz")
+
     def test_sets_submit_directory_for_s3_source_dir(self, tmp_path):
         """Should set sagemaker_submit_directory directly when source_dir is an S3 URI."""
         source_code = _make_source_code(source_dir="s3://my-bucket/code/sourcedir.tar.gz")