fix: return created dataset/evaluator from create_version, not bool

[papriwal]

Return created dataset/evaluator from create_version, not bool

Problem

DataSet.create_version and Evaluator.create_version both discarded the newly created version and returned only a boolean success flag (DataSet even silently swallowed all exceptions into False). Since create_version internally delegates to DataSet.create()/Evaluator.create() — which already return the full new entity — every caller who needed the new version's ARN was forced to make a separate get() or refresh() call afterward just to retrieve information the method already had in hand.

Changes

• DataSet.create_version now returns the newly created DataSet (or None on failure, matching its existing swallow-and-log pattern).
• Evaluator.create_version now returns the newly created Evaluator (still raises RuntimeError on failure, matching its existing behavior — no change to its error-handling contract).
• Both methods now thread sagemaker_session=self.sagemaker_session through to the delegated create() call. Previously this was omitted, so calling create_version on an entity created with a non-default session would silently create the new version under a different (default) session instead.
• Added a module-level logger = logging.getLogger(__name__) to dataset.py and evaluator.py (neither had one) and replaced the print() calls in create_version with logger.info/logger.error, matching the convention used elsewhere in sagemaker.train (e.g. sft_trainer.py, dpo_trainer.py).

Testing Done

Updated the existing unit tests in test_dataset.py and test_evaluator.py, and the integration tests in tests/integ/ai_registry/test_dataset.py and tests/integ/ai_registry/test_evaluator.py, to assert against the returned entity object instead of a boolean.

[github-actions]

Claude finished @papriwal's task in 1m 0s —— View job

---

PR Review: return created dataset/evaluator from create_version

• Read PR diff
• Verify create() signatures accept sagemaker_session
• Check callers of create_version for breakage
• Assess tests

Verdict: LGTM. This is a clean, well-motivated change and I found no correctness issues.

What I verified

• sagemaker_session threading is valid. Both DataSet.create (dataset.py:227) and Evaluator.create (evaluator.py:211) accept a sagemaker_session kwarg, and both entities store self.sagemaker_session (used elsewhere, e.g. evaluator.py:153, dataset.py:396). Passing it through fixes the real bug where a version was created under the default session. Good catch.
• self.type is populated (evaluator.py:128), so Evaluator.create(type=self.type, ...) is correct.
• No external callers break. create_version is only called from the updated tests — no other production callers rely on the old bool return. Existing truthiness checks (if x.create_version(...)) keep working since an entity is truthy and None is falsy.
• Error-handling contracts preserved. DataSet keeps its swallow-and-log-None pattern; Evaluator keeps raising RuntimeError. Both docstrings now document this accurately.
• Logging migration is consistent with the logger.info/logger.error %s-style convention used elsewhere in sagemaker.train.
• Tests now assert on the returned entity type/name and, for the unit tests, assert sagemaker_session is threaded into create() — which directly covers the session bug being fixed.

Minor observations (non-blocking)

1. Divergent failure contracts. DataSet.create_version returns None on failure while Evaluator.create_version raises. Both are pre-existing and intentionally preserved here, so no action needed — just flagging that a caller writing generic code against both would have to handle two different failure modes.
2. Public return-type change (bool → entity/None). Technically a signature change for SDK consumers, but the improvement clearly outweighs it and truthiness-based callers are unaffected. Worth a line in release notes if this API is considered stable/GA.

Nice, focused fix. 👍

[github-actions]

Claude finished @papriwal's task in 1m 3s —— View job

---

PR Review: return created dataset/evaluator from create_version

• Read PR diff
• Verify create() signatures & constructors
• Check callers of create_version
• Assess tests & error contracts
• Post review

Verdict: LGTM. Clean, well-motivated fix with no correctness issues found.

What I verified

• sagemaker_session threading is correct. Both DataSet.create (dataset.py:227) and Evaluator.create (evaluator.py:211) accept a sagemaker_session kwarg, and both entities store self.sagemaker_session via the base AIRHubEntity.__init__ (air_hub_entity.py:75). Passing it through fixes the real bug where the new version was silently created under the default session. Good catch.
• self.type is populated on the Evaluator instance, so Evaluator.create(type=self.type, ...) is valid.
• No callers break. create_version is only called from the updated tests and the example notebook (v3-examples/model-customization-examples/ai_registry_example.ipynb:161,322), which discards the return value — so the bool → entity/None change doesn't affect it. Any truthiness-based checks keep working (an entity is truthy, None is falsy).
• Error contracts preserved. DataSet keeps its swallow-and-log-None pattern; Evaluator keeps raising RuntimeError. Both docstrings now document this accurately.
• Logging migration is consistent with the logger.info/logger.error %s-style convention used elsewhere in sagemaker.train.
• Tests now assert on the returned entity type/name and, for the unit tests, assert sagemaker_session is threaded into create() — directly covering the session bug being fixed.

Minor observations (non-blocking)

1. Divergent failure contracts. DataSet.create_version returns None on failure while Evaluator.create_version raises RuntimeError. Both are pre-existing and intentionally preserved here — just flagging that generic code against both must handle two failure modes.
2. Public return-type change (bool → entity/None). Technically a signature change for SDK consumers; truthiness callers are unaffected, but worth a release-note line if this API is considered GA.

Nice, focused fix. 👍