Skip to content

Commit 3cbb262

Browse files
rohangujarathiRohan Gujarathi
authored andcommitted
change: Add a unit test for consistency between step and remote decorator (#1332)
Co-authored-by: Rohan Gujarathi <[email protected]>
1 parent 0faec72 commit 3cbb262

File tree

1 file changed

+35
-0
lines changed

1 file changed

+35
-0
lines changed

tests/unit/sagemaker/remote_function/test_client.py

+35
Original file line numberDiff line numberDiff line change
@@ -1425,3 +1425,38 @@ def test_list_future(mock_session):
14251425
call(TrainingJobName="job-3"),
14261426
]
14271427
)
1428+
1429+
1430+
def test_consistency_between_remote_and_step_decorator():
1431+
from sagemaker.workflow.function_step import step
1432+
1433+
remote_args_to_ignore = [
1434+
"_remote",
1435+
"include_local_workdir",
1436+
"custom_file_filter",
1437+
"s3_kms_key",
1438+
"s3_root_uri",
1439+
"sagemaker_session",
1440+
]
1441+
1442+
step_args_to_ignore = ["_step", "name", "display_name", "description", "retry_policies"]
1443+
1444+
remote_decorator_args = remote.__code__.co_varnames
1445+
common_remote_decorator_args = set(remote_args_to_ignore) ^ set(remote_decorator_args)
1446+
1447+
step_decorator_args = step.__code__.co_varnames
1448+
common_step_decorator_args = set(step_args_to_ignore) ^ set(step_decorator_args)
1449+
1450+
assert common_remote_decorator_args == common_step_decorator_args
1451+
1452+
1453+
def test_consistency_between_remote_and_executor():
1454+
executor_arg_list = list(RemoteExecutor.__init__.__code__.co_varnames)
1455+
executor_arg_list.remove("self")
1456+
executor_arg_list.remove("max_parallel_jobs")
1457+
1458+
remote_args_list = list(remote.__code__.co_varnames)
1459+
remote_args_list.remove("_remote")
1460+
remote_args_list.remove("_func")
1461+
1462+
assert executor_arg_list == remote_args_list

0 commit comments

Comments
 (0)