Instance type validation fix for serverful training jobs

[mujtaba1747]

Add instance_type validation for SMTJ training flow

Description

Adds instance_type validation to the training flow so that unsupported instance types are rejected client-side, before a training job is submitted.

SMTJ (SageMaker Training Job) hub content does not include an instance_type enum in its override spec, but the SMHP (SageMaker HyperPod) recipe for the same model/technique does. This change reuses the SMHP override spec to source the allowed instance types and applies them to SMTJ recipe validation, closing a gap where invalid instance types would only fail later at job launch.

Changes

• common_utils/finetune_utils.py: Add _get_smhp_instance_type_enum(...), which fetches the instance_type enum from the SMHP override spec for a given model/technique/training type. Returns None (with a warning) when the enum is unavailable, so validation degrades gracefully instead of failing.
• base_trainer.py:
  • Add _validate_instance_type(...) that resolves the allowed enum and raises a ValueError listing the supported values when the provided instance_type is not allowed.
  • Wire instance-type validation into the recipe-preparation path and emit a warning when the SMHP recipe does not provide a valid enum (validation skipped).
  • Add a matching warning for the existing instance-count validation when no valid replicas enum is available.
• Tests (integ):
  • test_sft_trainer_integration.py: assert an unsupported instance type (ml.t3.medium) raises ValueError for a Llama LORA serverful job.
  • test_sft_trainer_serverful_smtj.py: assert an unsupported instance type raises ValueError for a Nova SMTJ job, and that an out-of-range instance count raises ValueError when a valid instance type is used.

Behavior

• Invalid instance_type → ValueError: Instance type '<type>' is not supported. Allowed values: [...] raised during train(dry_run=True) before any job is submitted.
• When the SMHP recipe provides no enum, validation is skipped and a warning is logged (non-blocking).

Testing

• Added integration tests covering invalid instance type (serverful + SMTJ) and invalid instance count (SMTJ), each asserting the job fails fast with a clear ValueError.

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

[github-actions]

Claude finished @mujtaba1747's task in 2m 16s —— View job

---

PR Review — Instance type validation fix for serverful training jobs

• Read PR diff
• Review base_trainer.py changes
• Review finetune_utils.py changes
• Review notifications.py changes
• Review test changes

Overall the instance-type validation change is well structured and mirrors the existing _get_smhp_replicas_enum / _validate_instance_count pattern cleanly. A few things worth addressing, mostly in notifications.py.

Note: the inline-comment tool isn't available in this run, so findings are below with file/line references.

---

Correctness

1. notifications.py:151-165 — AWS:SourceAccount on the admin statement likely won't grant anything for direct IAM calls.
The SNSTopicAdministration statement (Principal: {"AWS": "*"}) is gated by Condition: {"StringEquals": {"AWS:SourceAccount": account_id}}. aws:SourceAccount is a global key that is only populated when an AWS service makes the request on your behalf — it is not present in the context of a direct Subscribe/GetTopicAttributes/SetTopicAttributes call by an IAM principal, so the Allow won't apply to those actions. AWS's default SNS topic policy uses the SNS-specific AWS:SourceOwner key for exactly this purpose. If the intent is to replicate the owner-admin default, this statement should likely use AWS:SourceOwner (the EventBridge statement at 154-161 correctly uses AWS:SourceAccount since a service populates it there). Worth verifying against how you tested this policy.

2. notifications.py:201-229 — create_notification_topic unconditionally overwrites the topic policy.
create_topic is idempotent and returns an existing topic's ARN if the name already exists; the subsequent set_topic_attributes then replaces whatever policy that topic had. If a caller reuses a name for a pre-existing topic with a custom policy, it's silently clobbered. Consider documenting this in the docstring, or only setting the policy on a freshly-created topic.

3. notifications.py:322-329 — behavior change: exception type and lost guidance.
The previous except Exception produced a ValueError that told the user exactly what resource-policy statement to add. The new code only maps AccessDenied → PermissionError and re-raises everything else raw. Two implications:

• Consumers previously catching ValueError here will now see PermissionError / raw ClientError — a small backward-incompatibility.
• The actionable "add this to your topic policy" message is gone. Since create_notification_topic now sets that policy, that's presumably intentional, but if a user brings their own topic they've lost the hint. Consider keeping a short pointer in the PermissionError message.

Minor

4. Redundant double warning on fetch failure. In base_trainer.py:913-919 (and the existing 930-935 replicas path), when _get_smhp_instance_type_enum hits an exception it already logs a warning internally (finetune_utils.py:100-104) and returns None; _train_serverful_smtj then logs a second, nearly identical warning. On a real fetch error the user sees two warnings. Not blocking, but you may want to log only in one place.

Tests

5. No unit coverage for the new logic. The two new files add only integ tests, which require live AWS resources and won't run in normal CI. The new pure/near-pure functions — _build_topic_policy (JSON shape/conditions), _get_smhp_instance_type_enum (enum present / missing / exception → None), and _validate_instance_type (raises vs. skips when enum is None) — are easy to unit-test with mocks and would guard against regressions like #1 above. Recommend adding unit tests.

---

Nothing here is a hard blocker except confirming the policy behavior in #1. The validation flow itself looks correct and consistent with existing conventions.