OFA PyTorch sample - Ensure full code locality

Author: eitansela

.gitignore

@@ -179,3 +179,4 @@ xgboost_script_mode_local_training_and_serving/code/build/lib/build/lib/build/li
 xgboost_script_mode_local_training_and_serving/model/model.tar.gz
 xgboost_script_mode_local_training_and_serving/model/output.tar.gz
 xgboost_script_mode_local_training_and_serving/model/xgboost-model
+pytorch_extend_dlc_container_ofa_local_serving/model.tar.gz

pytorch_extend_dlc_container_ofa_local_serving/pytorch_extend_dlc_container_ofa_local_serving.py

@@ -15,26 +15,34 @@
 from PIL import Image
 import numpy as np
 from sagemaker.pytorch import PyTorchModel
+from sagemaker.local import LocalSession
+import boto3
 
 DUMMY_IAM_ROLE = 'arn:aws:iam::111111111111:role/service-role/AmazonSageMaker-ExecutionRole-20200101T000001'
+LOCAL_SESSION = LocalSession()
+LOCAL_SESSION.config={'local': {'local_code': True}} # Ensure full code locality, see: https://sagemaker.readthedocs.io/en/stable/overview.html#local-mode
 
 
 def main():
 
  image = 'sagemaker-ofa-pytorch-extended-local'
 
+ print('Downloading model file from S3')
+ s3 = boto3.client('s3')
+ s3.download_file('aws-ml-blog', 'artifacts/pytorch-extend-dlc-container-ofa-tiny/model.tar.gz', 'model.tar.gz')
+ print('Model downloaded')
+
  ofa_hf_model = PyTorchModel(
  source_dir="code",
  entry_point="inference.py",
  role=DUMMY_IAM_ROLE,
- model_data="s3://aws-ml-blog/artifacts/pytorch-extend-dlc-container-ofa-tiny/model.tar.gz",
+ model_data="./model.tar.gz",
  image_uri=image,
- framework_version='1.8'
+ framework_version='1.8',
+ sagemaker_session=LOCAL_SESSION
  )
 
  print('Deploying endpoint in local mode')
- print(
- 'Note: model download might take a few minutes to complete due to its size.')
  predictor = ofa_hf_model.deploy(
  initial_instance_count=1,
  instance_type='local',