Make helper methods internal

Author: aviruthen

sagemaker-train/pyproject.toml

@@ -42,7 +42,6 @@ dependencies = [
     "jinja2>=3.0,<4.0",
     "sagemaker-mlflow>=0.0.1,<1.0.0",
     "mlflow>=3.0.0,<4.0.0",
-    "nest_asyncio>=1.5.0",
 ]
 
 [project.urls]

sagemaker-train/src/sagemaker/train/aws_batch/batch_api_helper.py

@@ -23,7 +23,7 @@
 from sagemaker.train.aws_batch.boto_client import get_batch_boto_client
 
 
-def submit_service_job(
+def _submit_service_job(
     training_payload: Dict,
     job_name: str,
     job_queue: str,
@@ -71,7 +71,7 @@ def submit_service_job(
     return client.submit_service_job(**payload)
 
 
-def describe_service_job(job_id: str) -> Dict:
+def _describe_service_job(job_id: str) -> Dict:
     """Batch describe_service_job API helper function.
 
     Args:
@@ -119,7 +119,7 @@ def describe_service_job(job_id: str) -> Dict:
     return client.describe_service_job(jobId=job_id)
 
 
-def terminate_service_job(job_id: str, reason: Optional[str] = "default terminate reason") -> Dict:
+def _terminate_service_job(job_id: str, reason: Optional[str] = "default terminate reason") -> Dict:
     """Batch terminate_service_job API helper function.
 
     Args:
@@ -132,7 +132,7 @@ def terminate_service_job(job_id: str, reason: Optional[str] = "default terminat
     return client.terminate_service_job(jobId=job_id, reason=reason)
 
 
-def list_service_job(
+def _list_service_job(
     job_queue: str,
     job_status: Optional[str] = None,
     filters: Optional[List] = None,
@@ -161,7 +161,7 @@ def list_service_job(
     next_token = part_of_jobs.get("nextToken")
     yield part_of_jobs
     if next_token:
-        yield from list_service_job(job_queue, job_status, filters, next_token)
+        yield from _list_service_job(job_queue, job_status, filters, next_token)
 
 
 def __merge_tags(batch_tags: Optional[Dict], training_tags: Optional[List]) -> Optional[Dict]:

sagemaker-train/src/sagemaker/train/aws_batch/training_queue.py

@@ -17,7 +17,7 @@
 import logging
 from sagemaker.train.model_trainer import ModelTrainer, Mode
 from .training_queued_job import TrainingQueuedJob
-from .batch_api_helper import submit_service_job, list_service_job
+from .batch_api_helper import _submit_service_job, _list_service_job
 from .exception import MissingRequiredArgument
 from .constants import DEFAULT_TIMEOUT, JOB_STATUS_RUNNING
 
@@ -85,7 +85,7 @@ def submit(
         if job_name is None:
             job_name = training_payload["TrainingJobName"]
 
-        resp = submit_service_job(
+        resp = _submit_service_job(
             training_payload,
             job_name,
             self.queue_name,
@@ -177,7 +177,7 @@ def list_jobs(
             status = None  # job_status is ignored when job_name is specified.
         jobs_to_return = []
         next_token = None
-        for job_result_dict in list_service_job(self.queue_name, status, filters, next_token):
+        for job_result_dict in _list_service_job(self.queue_name, status, filters, next_token):
             for job_result in job_result_dict.get("jobSummaryList", []):
                 if "jobArn" in job_result and "jobName" in job_result:
                     jobs_to_return.append(

sagemaker-train/src/sagemaker/train/aws_batch/training_queued_job.py

@@ -29,9 +29,9 @@
     SourceCode,
     TrainingImageConfig,
 )
-from .batch_api_helper import terminate_service_job, describe_service_job
+from .batch_api_helper import _terminate_service_job, _describe_service_job
 from .exception import NoTrainingJob, MissingRequiredArgument
-from ..utils import get_training_job_name_from_training_job_arn
+from ..utils import _get_training_job_name_from_training_job_arn
 from .constants import JOB_STATUS_COMPLETED, JOB_STATUS_FAILED, POLL_IN_SECONDS
 
 logging.basicConfig(
@@ -82,15 +82,15 @@ def terminate(self, reason: Optional[str] = "Default terminate reason") -> None:
         Returns: None
 
         """
-        terminate_service_job(self.job_arn, reason)
+        _terminate_service_job(self.job_arn, reason)
 
     def describe(self) -> Dict:
         """Describe Batch job.
 
         Returns: A dict which includes job parameters, job status, attempts and so on.
 
         """
-        return describe_service_job(self.job_arn)
+        return _describe_service_job(self.job_arn)
 
     def _training_job_created(self, status: str) -> bool:
         """Return True if a Training job has been created
@@ -320,7 +320,7 @@ def _get_new_training_job_name_from_latest_attempt(latest_attempt: Dict) -> str:
 
     """
     training_job_arn = latest_attempt.get("serviceResourceId", {}).get("value", None)
-    return get_training_job_name_from_training_job_arn(training_job_arn)
+    return _get_training_job_name_from_training_job_arn(training_job_arn)
 
 
 def _remove_system_tags_in_place_in_model_trainer_object(model_trainer: ModelTrainer) -> None:

sagemaker-train/src/sagemaker/train/utils.py

@@ -249,7 +249,7 @@ def _get_studio_tags(model_id: str, hub_name: str):
     ]
 
 
-def get_training_job_name_from_training_job_arn(training_job_arn: str) -> str:
+def _get_training_job_name_from_training_job_arn(training_job_arn: str) -> str:
     """Extract Training job name from Training job arn.
     Args:
         training_job_arn: Training job arn.