[egorkosaretsky]

Problem

Local mode endpoint deployment uses a hard-coded HEALTH_CHECK_TIMEOUT_LIMIT of 120s (entities.py:43). This causes failures when containers take longer to start — e.g. large model archives (~5GB+), slow pip installs, or low-bandwidth environments — even though the actual SageMaker endpoint deployment succeeds fine.

Fixes #3362.

Solution

• Added an optional timeout parameter to _wait_for_serving_container() (defaults to HEALTH_CHECK_TIMEOUT_LIMIT for backwards compatibility)
• Both _LocalEndpoint.serve() and _LocalTransformJob.start() now read local.health_check_timeout from the session config and pass it through

Usage:

sess = LocalSession()
sess.config = {'local': {'health_check_timeout': 600}}  # 10 minutes

Testing

Added unit tests covering:

• Custom timeout value is respected by _wait_for_serving_container
• health_check_timeout from session config is passed through in _LocalEndpoint.serve()
• Default behaviour (120s) is unchanged when config is not set

Checklist

• Unit tests added/updated
• Backwards compatible (default unchanged)
• Both endpoint and transform job code paths updated

[egorkosaretsky]

@nargokul Could you please take a look, this should fix #3362

[github-actions]

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

---

Code Review — Local mode health-check timeout

• Read the PR diff and changed files
• Reviewed changed code against callers and existing conventions
• Posted review

---

Thanks for the fix — this is a clean, well-scoped change that addresses #3362. It mirrors the existing local.serving_port config pattern, keeps the default at HEALTH_CHECK_TIMEOUT_LIMIT for backwards compatibility, and updates both the endpoint and transform-job code paths. Tests cover the new behavior well.

Correctness / backwards-compat: ✅ Default preserved, public signature extended with a keyword-only-by-convention arg, both call sites updated. No config schema needed updating (local config keys aren't validated against a schema).

A few minor, non-blocking observations:

1. or fallback masks a 0 value (low severity)

entities.py:341-344 and 639-642:

health_check_timeout = (
    get_config_value("local.health_check_timeout", self.local_session.config)
    or HEALTH_CHECK_TIMEOUT_LIMIT
)

If a user sets health_check_timeout: 0 it silently falls back to 120s rather than being honored. This is consistent with the existing serving_port ... or 8080 idiom and 0 is a meaningless timeout anyway, so this is fine to leave — just flagging the behavior.

2. No type coercion on the config value (low severity)

The value flows straight into i >= timeout. If someone sets it as a string (e.g. "600"), that comparison raises TypeError in Python 3. Again consistent with how serving_port is handled today, so not a blocker, but a small int(...) cast would make it more robust against config files that stringify values.

3. Misleading comment in test_wait_custom_timeout (nit)

test_entities.py:123:

# With timeout=10 and step=5, only 2 iterations before i(=10) >= timeout(=10)
assert mock_perform_request.call_count == 1

The assertion (== 1) is correct, but the comment says "2 iterations" — there are 2 loop iterations but only 1 _perform_request call (the second iteration raises before the request). Worth tweaking the comment so it matches the assert.

Optional: consider adding a short note about the new local.health_check_timeout config key to the local-mode docs so it's discoverable, since it's a user-facing option.

Nothing here blocks merge — the core logic and tests are solid. 👍

[mohamedzeidan2021]

Thanks for this — a couple of things to sort out before it can go in.

The config key isn't accepted by the local-mode schema. SAGEMAKER_PYTHON_SDK_LOCAL_MODE_CONFIG_SCHEMA sets additionalProperties: False and lists only local_code, region_name, serving_port, container_root, and container_config (config_schema.py:1195). LocalSession.config's setter validates against it (local_session.py:678), so setting the new key in a local-mode config file raises instead of extending the timeout:

jsonschema.validate({"local": {"health_check_timeout": 600}},
                    SAGEMAKER_PYTHON_SDK_LOCAL_MODE_CONFIG_SCHEMA)
# ValidationError: Additional properties are not allowed ('health_check_timeout' was unexpected)

The new tests don't catch this because they assign mock_session.config = {...} onto a Mock(), which bypasses validation. Could you add a health_check_timeout integer property to that schema, plus a case in tests/unit/config/test_config_schema.py?

Two smaller test notes:

• test_wait_uses_default_timeout_constant doesn't verify the default — it makes the first ping succeed and asserts one call, which passes for any default value. Making every ping fail and asserting the implied call count (120 / 5 = 23) would actually pin it.
• The _LocalTransformJob.start() hunk has no direct test; only the endpoint path is covered.

On CI: none of the unit or integ failures come from this change. unit-tests (sagemaker-core) and unit-tests (sagemaker-train) are branch staleness — the branch is 16 commits behind and has 1a3cecf5 (adds PyTorch 2.14 to the image config) but not 2e2ffc0e (the test fix that derives those versions from config) or 593aaadc (the HOME fix); merging master clears both. The integ failures are account quota and a stuck pipeline step.

codestyle-doc-tests (sagemaker-core) is the exception — it's partly this change. tests/unit/local/test_entities.py:765 exceeds 100 columns and black wants to rewrap it. The other 5 files in that job are pre-existing on master, so running black on the one file here won't turn the job green on its own, but it clears this PR's contribution to it.