Spaces:
Running on Zero
Running on Zero
File size: 4,750 Bytes
2e1dc7f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | """
This is the script that runs on the "controller" machine: the lightweight
machine that requests the big boi instance(s) from EC2 and configures the
training job. In this example, I want the controller machine to be my home
computer. This means that I can run this script from my computer and start a
training job on the remote big boi instance, without an intermediary machine
from aws.
This script can run in two modes, local and remote, depending on the value of
`run_locally`.
VERY IMPORTANT NOTE: running this script locally is not at all the same as
running a normal training locally (i.e: running locally the file `train.py`):
Running this script in local mode still creates a sagemaker job and does
everything in the same exact way as it would to run on a remote instance,
EXCEPT that the big boi instance in EC2 is replaced by your local machine.
So the script creates a docker container, downloads it, creates the virtual
environment, and runs the training job in the same exact way as it would on the
remote instance. Running this training job in local mode successfuly should
GUARANTEE that it will run successfully on a remote instance.
"""
import sagemaker.pytorch
import sagemaker.local
import sagemaker
from datetime import datetime
def launch_aws_job(
experiment_name: str,
instance_type: str,
max_training_hours: int,
entry_point: str,
input_mode: str = "File",
data_dir: str = "",
):
# True: run training job locally for testing purposes (needs docker installed)
# False: run training job on an AWS instance
run_locally: bool = False
# define names for project, experiment and run
project_name = "latent-accompaniment-generation"
# experiment_name = "bart-drums-mixdata"
checkpoint_ec2_path = "/opt/ml/checkpoints"
# checkpoint_s3_path = "s3://latent-accompaniment-generation/checkpoints/ash-drums"
# NOTE: job name has to be unique
now = datetime.now()
datestring = now.strftime('%Y-%m-%d-%H-%M')
job_name = experiment_name + "-" + datestring
# define session and role
sagemaker_session = (sagemaker.local.LocalSession()
if run_locally else sagemaker.Session())
role = ("arn:aws:iam::076456026604:role/service-role/"
"AmazonSageMaker-ExecutionRole-20240407T195108")
# define input and output directories
if run_locally:
data_path: str = "file://./data/mtg-jamendo-low/"
output_path: str = ("file://./checkpoints/"
f"{project_name}/{experiment_name}/{job_name}")
else:
# NOTE: whatever directory is specified in here will be downloaded to the
# training machine in its entirety (if using "File" input mode, which is
# the default). So make sure that the s3 directory specified here only
# contains data that is needed for the training.
# WRONG: s3://audio-data-bucket/data/ (.../MNIST, .../moisesdb)
# OK: s3://audio-data-bucket/data/MNIST
data_path: str = "s3://audio-data-bucket/data/lag-data/" + data_dir
weights_path: str = "s3://latent-accompaniment-generation/weights/"
output_path = f"s3://{project_name}/{experiment_name}"
# ckp_path = checkpoint_s3_path
# instance_type: str = "ml.g4dn.xlarge"
# instance_type: str = "ml.m5.xlarge"
# instance_type: str = "ml.g5.2xlarge" # 1x nvidia A10
# instance_type: str = "ml.g5.12xlarge" # 4x nvidia A10
# instance_type: str = "ml.p4d.24xlarge" # 8x nvidia A100 40gb 💀
# instance_type: str = "ml.p4de.24xlarge" # 8x nvidia A100 80gb 💀
# instance_type: str = "ml.trn1.2xlarge"
# max_training_time = 3 * 24 * 60 * 60 # 3 days
# max_training_time = 16 * 60 * 60 # 16 hours
max_training_time = max_training_hours * 60 * 60
estimator = sagemaker.pytorch.PyTorch(
# entry_point="scripts/train_bart.py",
entry_point=entry_point,
role=role,
max_run=max_training_time,
instance_count=1,
framework_version="2.1.0",
py_version="py310",
dependencies=["requirements.txt"],
source_dir="src",
output_path=output_path,
instance_type="local" if run_locally else instance_type,
local_code=run_locally,
# input_mode="FastFile",
# input_mode="File",
input_mode=input_mode,
checkpoint_local_path=checkpoint_ec2_path,
# checkpoint_s3_uri=checkpoint_s3_path,
# distribution={"pytorchddp": {
# "enabled": "true"
# }},
)
estimator.fit(
{
"data": data_path,
# "ckp": ckp_path,
"weights": weights_path,
},
job_name=job_name,
)
|