Support attaching to already submitted training jobs with SSHEstimatorWrapper.attach(job_name)

Author: ivan-khvostishkov

FAQ.md

@@ -153,10 +153,27 @@ You might want (optionally) to configure [AWS PrivateLink for Session Manager en
 ## API Questions
 
 ### I'm using boto3 Python SDK instead of SageMaker Python SDK, how can I use SageMaker SSH Helper?
-This is a tricky question. In short, this use case is not supported by SageMaker SSH Helper.
+This use case is not fully supported by SageMaker SSH Helper.
 However, [you can](https://repost.aws/questions/QU8-U_XgPVRSuLTSXf8eW8fA/can-we-connect-to-the-instance-via-ssh-or-other-means-where-a-triton-sagemaker-endpoint-is-deployed) analyze the source code and re-implement SageMaker SSH Helper behaviour with boto3, e.g., by passing environment variables from your code.
+
 In general, this is not recommended, because the set of environment variables and internal logic is a subject to future changes. These changes won't necessarily appear in the release notes and can break your code.
 
+However, if submitted a job with boto3 in this way and started SSH Helper inside the container, you can use high-level APIs to fetch instance IDs and connect to the job with `sm-local` scripts from your local machine: 
+
+```python
+import logging
+from sagemaker_ssh_helper.wrapper import SSHEstimatorWrapper
+
+training_job_name = ...
+
+ssh_wrapper = SSHEstimatorWrapper.attach(training_job_name)
+
+instance_ids = ssh_wrapper.get_instance_ids()
+
+logging.info(f"To connect over SSM run: aws ssm start-session --target {instance_ids[0]}")
+logging.info(f"To connect over SSH run: sm-local-ssh-training connect {ssh_wrapper.latest_training_job_name()}")
+```
+
 
 ### How can I change the SSH authorized keys bucket and location when running `sm-local-ssh-*` commands?
 The **public** key is transferred to the container through the default SageMaker bucket with the S3 URI that looks 

sagemaker_ssh_helper/detached_sagemaker.py

@@ -0,0 +1,31 @@
+import sagemaker
+from sagemaker import Session
+# noinspection PyProtectedMember
+from sagemaker.estimator import _TrainingJob
+
+
+class DetachedEstimator(sagemaker.estimator.EstimatorBase):
+    f"""
+    A sagemaker.estimator.Estimator that does not block on attach().
+    """
+
+    def __init__(self, training_job_name: str, sagemaker_session: Session):
+        super().__init__(sagemaker_session=sagemaker_session, instance_count=0)
+        self._current_job_name = training_job_name
+        self.latest_training_job = _TrainingJob(sagemaker_session, self._current_job_name)
+
+    def training_image_uri(self):
+        raise ValueError("Not implemented")
+
+    def hyperparameters(self):
+        raise ValueError("Not implemented")
+
+    def create_model(self, **kwargs):
+        raise ValueError("Not implemented")
+
+    @classmethod
+    def attach(cls, training_job_name: str, sagemaker_session=None, model_channel_name="model"):
+        if not isinstance(training_job_name, str):
+            raise ValueError("training_job_name MUST be a string")
+        # TODO: fetch job details and call _prepare_init_params_from_job_description()
+        return DetachedEstimator(training_job_name, sagemaker_session)

tests/test_attach.py

@@ -0,0 +1,60 @@
+import logging
+import os
+from datetime import timedelta
+
+import pytest
+from sagemaker import Session
+# noinspection PyProtectedMember
+from sagemaker.estimator import _TrainingJob
+from sagemaker.pytorch import PyTorch
+
+from sagemaker_ssh_helper.detached_sagemaker import DetachedEstimator
+from sagemaker_ssh_helper.wrapper import SSHEstimatorWrapper
+
+
+# noinspection DuplicatedCode
+def test_attach_estimator():
+    estimator = PyTorch(entry_point=os.path.basename('source_dir/training/train.py'),
+                        source_dir='source_dir/training/',
+                        dependencies=[SSHEstimatorWrapper.dependency_dir()],
+                        base_job_name='ssh-training',
+                        framework_version='1.9.1',
+                        py_version='py38',
+                        instance_count=1,
+                        instance_type='ml.m5.xlarge',
+                        max_run=int(timedelta(minutes=15).total_seconds()),
+                        keep_alive_period_in_seconds=1800,
+                        container_log_level=logging.INFO)
+
+    _ = SSHEstimatorWrapper.create(estimator, connection_wait_time_seconds=600)
+
+    estimator.fit(wait=False)
+
+    job: _TrainingJob = estimator.latest_training_job
+    ssh_wrapper = SSHEstimatorWrapper.attach(job.name)
+
+    instance_ids = ssh_wrapper.get_instance_ids()
+
+    logging.info(f"To connect over SSM run: aws ssm start-session --target {instance_ids[0]}")
+    logging.info(f"To connect over SSH run: sm-local-ssh-training connect {ssh_wrapper.latest_training_job_name()}")
+
+    ssh_wrapper.start_ssm_connection_and_continue(11022, 60)
+
+    ssh_wrapper.wait_training_job()
+
+    assert estimator.model_data.find("model.tar.gz") != -1
+
+
+def test_cannot_fit_detached_estimator():
+    estimator = DetachedEstimator.attach('training-job-name', Session())
+
+    with pytest.raises(ValueError):
+        _ = SSHEstimatorWrapper.create(estimator)
+
+
+def test_can_fetch_job_name_from_detached_estimator():
+    ssh_wrapper = SSHEstimatorWrapper.attach('training-job-name', Session())
+
+    job_name = ssh_wrapper.latest_training_job_name()
+
+    assert job_name == 'training-job-name'