Write output to s3

sdgym/benchmark.py

@@ -281,6 +281,25 @@ def _run_on_dask(jobs, verbose):
     return dask.compute(*persisted)
 
 
+def _write_output_to_csv(scores, output_path, aws_key, aws_secret):
+    if not output_path:
+        return
+
+    if is_s3_path(output_path):
+        s3 = get_s3_client(aws_key, aws_secret)
+        bucket_name, key = parse_s3_path(output_path)
+        scores_csv = scores.to_csv(index=False)
+        s3.put_object(
+            Bucket=bucket_name,
+            Key=key,
+            Body=scores_csv.encode('utf-8'),
+        )
+    else:
+        scores.to_csv(output_path, index=False)
+
+    return
+
+
 def run(synthesizers, datasets=None, datasets_path=None, modalities=None, bucket=None,
         metrics=None, iterations=1, workers=1, cache_dir=None, show_progress=False,
         timeout=None, output_path=None, aws_key=None, aws_secret=None):
@@ -392,7 +411,7 @@ def run(synthesizers, datasets=None, datasets_path=None, modalities=None, bucket
         raise SDGymError("No valid Dataset/Synthesizer combination given")
 
     scores = pd.concat(scores)
-    if output_path:
-        scores.to_csv(output_path, index=False)
+
+    _write_output_to_csv(scores, output_path, aws_key, aws_secret)
 
     return scores