Cloud utilities refactor (#85)

* Revert back pytorch for better compatibility

* allow to use multiple retries

* fix AWS batch job name error

cleanrl_utils/submit_exp.py

@@ -6,6 +6,7 @@
 
 import boto3
 import requests
+import wandb
 
 # fmt: off
 parser = argparse.ArgumentParser(description='CleanRL Experiment Submission')
@@ -42,6 +43,8 @@
     help='if toggled, the script will push the built container')
 parser.add_argument('--provider', type=str, default="", choices=["aws"],
     help='the cloud provider of choice (currently only `aws` is supported)')
+parser.add_argument('--aws-num-retries', type=int, default=1,
+    help='the number of job retries for `provider=="aws"`')
 args = parser.parse_args()
 # fmt: on
 
@@ -87,7 +90,7 @@
 if args.provider == "aws":
     client = boto3.client("batch")
     for final_run_cmd in final_run_cmds:
-        job_name = args.command.replace(".py", "").replace("/", "_").replace(" ", "").replace("-", "_") + str(
+        job_name = args.docker_tag.replace(":", "").replace("/", "_").replace(" ", "").replace("-", "_") + str(
             int(time.time())
         )
         resources_requirements = []
@@ -119,10 +122,12 @@
                     "command": ["/bin/bash", "-c", final_run_cmd],
                     "environment": [
                         {"name": "WANDB_API_KEY", "value": args.wandb_key},
+                        {'name': 'WANDB_RESUME', 'value': 'allow'},
+                        {'name': 'WANDB_RUN_ID', 'value': wandb.util.generate_id()},
                     ],
                     "resourceRequirements": resources_requirements,
                 },
-                retryStrategy={"attempts": 1},
+                retryStrategy={"attempts": args.aws_num_retries},
                 timeout={"attemptDurationSeconds": int(args.num_hours * 60 * 60)},
             )
             if response["ResponseMetadata"]["HTTPStatusCode"] != 200: