|
25 | 25 | from unittest.mock import patch, MagicMock, ANY, mock_open |
26 | 26 | from tempfile import NamedTemporaryFile |
27 | 27 |
|
28 | | -from sagemaker.core.resources import TrainingJob |
| 28 | +from sagemaker.core.resources import Base, TrainingJob |
29 | 29 | from sagemaker.core.shapes import ( |
30 | 30 | ResourceConfig, |
31 | 31 | VpcConfig, |
|
45 | 45 | ) |
46 | 46 | from sagemaker.core.config.config_manager import SageMakerConfig |
47 | 47 | from sagemaker.core.helper.session_helper import Session |
| 48 | +from sagemaker.core.workflow.pipeline_context import PipelineSession |
48 | 49 | from sagemaker.train.model_trainer import ModelTrainer, Mode |
49 | 50 | from sagemaker.train.defaults import DEFAULT_INSTANCE_TYPE |
50 | 51 | from sagemaker.train.constants import ( |
@@ -2616,3 +2617,106 @@ def test_output_data_strips_trailing_slash_on_fallback(model_trainer): |
2616 | 2617 | assert model_trainer.output_data == ( |
2617 | 2618 | f"s3://{DEFAULT_BUCKET}/{DEFAULT_BUCKET_PREFIX}" "/my-training-job/output/output.tar.gz" |
2618 | 2619 | ) |
| 2620 | + |
| 2621 | + |
| 2622 | +# --------------------------------------------------------------------------- |
| 2623 | +# Tag serialization (issue #6474) |
| 2624 | +# --------------------------------------------------------------------------- |
| 2625 | + |
| 2626 | +TAG = Tag(key="project", value="ml-project") |
| 2627 | +BOTO_TAG = {"Key": "project", "Value": "ml-project"} |
| 2628 | + |
| 2629 | + |
| 2630 | +def _pipeline_session(): |
| 2631 | + """A real PipelineSession with only its outbound edges stubbed.""" |
| 2632 | + session = PipelineSession.__new__(PipelineSession) |
| 2633 | + session._context = None |
| 2634 | + session.boto_session = MagicMock() |
| 2635 | + session.boto_session.region_name = DEFAULT_REGION |
| 2636 | + session.sagemaker_config = {} |
| 2637 | + session.default_bucket = lambda: DEFAULT_BUCKET |
| 2638 | + session.default_bucket_prefix = DEFAULT_BUCKET_PREFIX |
| 2639 | + return session |
| 2640 | + |
| 2641 | + |
| 2642 | +def _pipeline_trainer(tags): |
| 2643 | + """A trainer whose session makes ``_create_training_job_args`` serialize the request.""" |
| 2644 | + return ModelTrainer( |
| 2645 | + training_image=DEFAULT_IMAGE, |
| 2646 | + role=DEFAULT_ROLE, |
| 2647 | + sagemaker_session=_pipeline_session(), |
| 2648 | + compute=DEFAULT_COMPUTE_CONFIG, |
| 2649 | + stopping_condition=DEFAULT_STOPPING_CONDITION, |
| 2650 | + output_data_config=DEFAULT_OUTPUT_DATA_CONFIG, |
| 2651 | + tags=tags, |
| 2652 | + ) |
| 2653 | + |
| 2654 | + |
| 2655 | +def _train_wire_tags(trainer): |
| 2656 | + """The Tags CreateTrainingJob really receives, measured not derived.""" |
| 2657 | + mock_client = MagicMock() |
| 2658 | + with patch.object(Base, "get_sagemaker_client", return_value=mock_client): |
| 2659 | + try: |
| 2660 | + trainer.train(wait=False) |
| 2661 | + except ValidationError: |
| 2662 | + # Building the TrainingJob resource out of a MagicMock response fails after |
| 2663 | + # the call. The call itself is what is under test. |
| 2664 | + pass |
| 2665 | + assert ( |
| 2666 | + mock_client.create_training_job.call_args is not None |
| 2667 | + ), "create_training_job was not reached" |
| 2668 | + return mock_client.create_training_job.call_args.kwargs.get("Tags") |
| 2669 | + |
| 2670 | + |
| 2671 | +def test_pipeline_request_uses_boto_style_tag_keys(): |
| 2672 | + """Regression for #6474. |
| 2673 | +
|
| 2674 | + The service requires Tags of the form {"Key": ..., "Value": ...}. ``serialize`` |
| 2675 | + renames the fields of Tag models but passes plain dicts through untouched, so |
| 2676 | + dumping the tags to dicts first left lowercase keys in the pipeline definition |
| 2677 | + and CreatePipeline rejected it with "Unknown Argument member 'key'". |
| 2678 | + """ |
| 2679 | + args = _pipeline_trainer([TAG])._create_training_job_args() |
| 2680 | + |
| 2681 | + assert args["Tags"] == [BOTO_TAG] |
| 2682 | + |
| 2683 | + |
| 2684 | +def test_pipeline_request_omits_tags_when_unset(): |
| 2685 | + """An untagged trainer must not add an empty Tags member to the request.""" |
| 2686 | + args = _pipeline_trainer(None)._create_training_job_args() |
| 2687 | + |
| 2688 | + assert "Tags" not in args |
| 2689 | + |
| 2690 | + |
| 2691 | +def test_boto3_request_uses_boto_style_tag_keys(model_trainer): |
| 2692 | + """The boto3 payload (used by TrainingQueue) needs the same boto-style tags.""" |
| 2693 | + model_trainer.tags = [TAG] |
| 2694 | + |
| 2695 | + args = model_trainer._create_training_job_args(boto3=True) |
| 2696 | + |
| 2697 | + assert args["Tags"] == [BOTO_TAG] |
| 2698 | + |
| 2699 | + |
| 2700 | +def test_interactive_request_keeps_tags_training_job_create_accepts(model_trainer): |
| 2701 | + """The interactive request carries dicts, not Tag models. |
| 2702 | +
|
| 2703 | + ``TrainingJob.create`` annotates ``tags`` with ``sagemaker.core.resources.Tag``, |
| 2704 | + which shadows the shape of the same name, so a ``sagemaker.core.shapes.Tag`` |
| 2705 | + instance fails its validation. The dicts are revalidated into that model instead, |
| 2706 | + and come out boto-style on the wire. |
| 2707 | + """ |
| 2708 | + model_trainer.tags = [TAG] |
| 2709 | + |
| 2710 | + args = model_trainer._create_training_job_args() |
| 2711 | + |
| 2712 | + assert args["tags"] == [{"key": "project", "value": "ml-project"}] |
| 2713 | + assert _train_wire_tags(model_trainer) == [BOTO_TAG] |
| 2714 | + |
| 2715 | + |
| 2716 | +def test_pipeline_and_interactive_emit_identical_tags(model_trainer): |
| 2717 | + """Both paths must put the same tags on the request.""" |
| 2718 | + model_trainer.tags = [TAG] |
| 2719 | + |
| 2720 | + pipeline_tags = _pipeline_trainer([TAG])._create_training_job_args()["Tags"] |
| 2721 | + |
| 2722 | + assert pipeline_tags == _train_wire_tags(model_trainer) |
0 commit comments