feat: include training_plan_arn in remote function classes (#5157)

[mohamedzeidan2021]

Problem

The remote function interfaces (@remote decorator, RemoteExecutor, and the underlying _JobSettings) did not expose training_plan_arn, so remote-function training jobs could not target reserved-capacity training plans. ResourceConfig.TrainingPlanArn is a valid member of CreateTrainingJob in botocore (no service dependency), but there was no way to set it.

Fix

Added an optional training_plan_arn: str = None parameter to @remote, RemoteExecutor.__init__, and _JobSettings.__init__ (sagemaker-core). It is forwarded into _JobSettings from both public entry points, stored on the settings, and — when set — added to the compiled ResourceConfig as TrainingPlanArn using the same is not None guard as the adjacent optional fields. The parameter defaults to None, so behavior is byte-for-byte unchanged when omitted. Docstrings added in all three locations. (The sagemaker-train remote_function modules are deprecation shims that re-export from core, so they inherit the change automatically.)

Testing

• test_compile_with_training_plan_arn / test_compile_without_training_plan_arn: TrainingPlanArn present in ResourceConfig when set, absent when None.
• test_init_stores_training_plan_arn / test_init_training_plan_arn_defaults_to_none: real _JobSettings stores the value / defaults to None.
• test_remote_decorator_forwards_training_plan_arn / test_remote_executor_forwards_training_plan_arn: both public entry points forward the kwarg into _JobSettings.
• Full remote_function unit suite: 236 passed, no regressions. black + flake8 clean.

Fixes #5157

[github-actions]

Claude finished @mohamedzeidan2021's task in 55s —— View job

---

Code Review — training_plan_arn in remote function classes

• Read the PR diff
• Inspect changed source files & surrounding context
• Verify training_plan_arn typing & conventions
• Check test coverage
• Post review

This is a clean, well-scoped change. The new parameter is threaded consistently through all three entry points (@remote, RemoteExecutor.__init__, _JobSettings.__init__), stored on the settings, and injected into ResourceConfig using the same is not None guard as the adjacent optional fields (VolumeKmsKeyId, KeepAlivePeriodInSeconds). Behavior is unchanged when the arg is omitted, and the test coverage is thorough — both entry points, the storage/default on _JobSettings, and the presence/absence in the compiled ResourceConfig. Nice work.

Suggestion (minor)

Type annotation should be Optional[str] — client.py:102, client.py:567, job.py:422

training_plan_arn: str = None,

The default is None but the annotation is bare str, which is technically incorrect and inconsistent with the adjacent new-style optionals in these same signatures (e.g. nproc_per_node: Optional[int] = None). Optional is already imported in both modules. Prefer:

training_plan_arn: Optional[str] = None,

Purely a typing/consistency polish — no functional impact.

Notes / things I checked and consider fine

• No SDK-config resolution (resolve_value_from_config) for training_plan_arn. That's consistent with the neighboring instance_count / volume_size / keep_alive_period_in_seconds, which are also assigned directly, so this matches the surrounding pattern.
• No input validation of the ARN — appropriate here; the value is forwarded to botocore/CreateTrainingJob and validated service-side, same as the other pass-through fields.
• Deprecation shims in sagemaker-train re-export from core, so they inherit the change automatically, as the PR description states.
• No security or performance concerns; no backward-incompatibility (additive optional kwarg defaulting to None).

No blocking issues — only the Optional[str] annotation nit above.

[github-actions]

Claude finished @mohamedzeidan2021's task in 44s —— View job

---

Code Review — training_plan_arn in remote function classes

• Read the PR diff (/tmp/pr.diff)
• Inspect changed source & surrounding context (client.py, job.py, field-consistency guard)
• Verify correctness, conventions, and test coverage
• Post review

---

This is a clean, tightly-scoped additive change. training_plan_arn is threaded consistently through all three entry points (@remote, RemoteExecutor.__init__, _JobSettings.__init__), stored on the settings, and injected into ResourceConfig["TrainingPlanArn"] guarded by is not None — matching the adjacent optional fields (VolumeKmsKeyId, KeepAlivePeriodInSeconds). TrainingPlanArn is a valid member of ResourceConfig in the CreateTrainingJob API, so the placement is correct. Behavior is byte-for-byte unchanged when the argument is omitted.

What I verified:

• Correctness — the compile injection sits alongside the other optional ResourceConfig keys and uses the same guard; no ordering/merge concerns.
• Field-consistency guard — test_feature_scheduler.py::test_remote_decorator_fields_consistency asserts exact set equality between _JobSettings.__dict__ keys and an expected set. Since the change adds exactly one new self. attribute, adding "training_plan_arn" to that set keeps the guard passing. Good that this wasn't missed.
• Test coverage — thorough: both public entry points forward the kwarg (test_client.py), _JobSettings stores it / defaults to None (test_job.py), and presence/absence in the compiled ResourceConfig is covered.
• Deprecation shims — the sagemaker-train remote_function modules re-export from core, so they inherit the change automatically, as the description states.
• Security / perf / compatibility — no concerns. Value is a pass-through to CreateTrainingJob, validated service-side (consistent with the neighboring pass-through fields); purely additive optional kwarg.

Minor (non-blocking) — type annotation consistency

client.py:102, client.py:567, job.py:422:

training_plan_arn: str = None,

The default is None but the annotation is bare str. For consistency with the adjacent new-style optionals in these same signatures (e.g. nproc_per_node: Optional[int] = None), prefer Optional[str] = None. Optional is already imported in both modules. Purely a typing polish — no functional impact.

No blocking issues.
· fix/issue-5157-remote-training-plan-arn

[lucasjia-aws]

Automated review: 1 non-blocking suggestion. See inline comment.

[lucasjia-aws]

[suggestion] typing

The annotation str contradicts the None default — a type checker reads training_plan_arn: str = None as an invalid default. The adjacent optional parameter uses Optional[int] = None; match that here so the annotation is honest about accepting None.

training_plan_arn: Optional[str] = None,

Same pattern at client.py:567 (RemoteExecutor.__init__) and job.py:422 (_JobSettings.__init__).