Merge pull request #208 from Paperspace/PS-12571-Support_datasets_storage_configuration_options

Add support for experiment storage options. Removed some duplicated t…

Author: BartoszCki

gradient/api_sdk/clients/experiment_client.py

@@ -2,7 +2,7 @@
 
 from .base_client import BaseClient
 from .. import repositories, models, constants, utils
-from ..sdk_exceptions import InvalidParametersError
+from ..sdk_exceptions import ResourceCreatingDataError, InvalidParametersError
 from ..validation_messages import EXPERIMENT_MODEL_PATH_VALIDATION_ERROR
 
 
@@ -72,6 +72,8 @@ def create_single_node(
         if not is_preemptible:
             is_preemptible = None
 
+        datasets = self._dataset_dicts_to_instances(datasets)
+
         experiment = models.SingleNodeExperiment(
             experiment_type_id=constants.ExperimentType.SINGLE_NODE,
             name=name,
@@ -191,6 +193,8 @@ def create_multi_node(
         if not is_preemptible:
             is_preemptible = None
 
+        datasets = self._dataset_dicts_to_instances(datasets)
+
         experiment = models.MultiNodeExperiment(
             name=name,
             project_id=project_id,
@@ -314,6 +318,7 @@ def create_mpi_multi_node(
         if not is_preemptible:
             is_preemptible = None
 
+        datasets = self._dataset_dicts_to_instances(datasets)
         experiment_type_id = constants.ExperimentType.MPI_MULTI_NODE
 
         experiment = models.MpiMultiNodeExperiment(
@@ -421,6 +426,8 @@ def run_single_node(
         if not is_preemptible:
             is_preemptible = None
 
+        datasets = self._dataset_dicts_to_instances(datasets)
+
         experiment = models.SingleNodeExperiment(
             experiment_type_id=constants.ExperimentType.SINGLE_NODE,
             name=name,
@@ -538,6 +545,8 @@ def run_multi_node(
         if not is_preemptible:
             is_preemptible = None
 
+        datasets = self._dataset_dicts_to_instances(datasets)
+
         experiment = models.MultiNodeExperiment(
             name=name,
             project_id=project_id,
@@ -661,6 +670,8 @@ def run_mpi_multi_node(
         if not is_preemptible:
             is_preemptible = None
 
+        datasets = self._dataset_dicts_to_instances(datasets)
+
         experiment_type_id = constants.ExperimentType.MPI_MULTI_NODE
 
         experiment = models.MpiMultiNodeExperiment(
@@ -840,3 +851,23 @@ def stream_metrics(self, experiment_id, interval="30s", built_in_metrics=None):
             built_in_metrics=built_in_metrics,
         )
         return metrics
+
+    def _dataset_dicts_to_instances(self, datasets):
+        if not datasets:
+            return None
+
+        if isinstance(datasets, dict):
+            datasets = [datasets]
+
+        for ds in datasets:
+            if not ds.get("uri"):
+                raise ResourceCreatingDataError("Error while creating experiment with dataset: "
+                                                "\"uri\" key is required and it's value must be a valid S3 URI")
+
+        for ds in datasets:
+            volume_options = ds.setdefault("volume_options", {})
+            volume_options.setdefault("kind", ds.pop("volume_kind", None))
+            volume_options.setdefault("size", ds.pop("volume_size", None))
+
+        datasets = [models.Dataset(**ds) for ds in datasets]
+        return datasets

gradient/api_sdk/constants.py

@@ -136,3 +136,16 @@ class BuiltinMetrics:
         ("gpuMemoryUtilization", BuiltinMetrics.gpu_memory_utilization),
     )
 )
+
+
+class DatasetVolumeKinds:
+    DYNAMIC = "dynamic"
+    SHARED = "shared"
+
+
+DATASET_VOLUME_KINDS = collections.OrderedDict(
+    (
+        ("dynamic", DatasetVolumeKinds.DYNAMIC),
+        ("shared", DatasetVolumeKinds.SHARED),
+    ),
+)

gradient/api_sdk/models/dataset.py

@@ -1,6 +1,12 @@
 import attr
 
 
+@attr.s
+class VolumeOptions(object):
+    kind = attr.ib(type=str, default=None)
+    size = attr.ib(type=str, default=None)
+
+
 @attr.s
 class Dataset(object):
     uri = attr.ib(type=str, default=None)
@@ -9,3 +15,4 @@ class Dataset(object):
     etag = attr.ib(type=str, default=None)
     version_id = attr.ib(type=str, default=None)
     name = attr.ib(type=str, default=None)
+    volume_options = attr.ib(type=VolumeOptions, factory=VolumeOptions)

gradient/api_sdk/serializers/dataset.py

@@ -5,6 +5,11 @@
 from .. import models
 
 
+class DatasetVolumeOptionsSchema(BaseSchema):
+    kind = ma.fields.Str(required=True)
+    size = ma.fields.Str(required=True)
+
+
 class DatasetSchema(BaseSchema):
     MODEL = models.Dataset
 
@@ -15,3 +20,4 @@ class DatasetSchema(BaseSchema):
     version_id = ma.fields.String(dump_to="versionId", load_from="versionId")
     name = ma.fields.String()
     tags = ma.fields.Nested(TagSchema, only="name", many=True, load_only=True)
+    volume_options = ma.fields.Nested(DatasetVolumeOptionsSchema, dump_to="volumeOptions", load_from="volumeOptions")

gradient/cli/experiments.py

@@ -225,6 +225,21 @@ def dataset_options(f):
             help="S3 dataset's ETag",
             cls=common.GradientDatasetOption,
         ),
+        click.option(
+            "--datasetVolumeKind",
+            "dataset_volume_kind_list",
+            multiple=True,
+            type=ChoiceType(constants.DATASET_VOLUME_KINDS, case_sensitive=False),
+            help="S3 dataset's volume kind. If used, --datasetVolumeSize has to be set as well",
+            cls=common.GradientDatasetOption,
+        ),
+        click.option(
+            "--datasetVolumeSize",
+            "dataset_volume_size_list",
+            multiple=True,
+            help="S3 dataset's volume size",
+            cls=common.GradientDatasetOption,
+        ),
     ]
     return functools.reduce(lambda x, opt: opt(x), reversed(options), f)
 

gradient/commands/experiments.py

@@ -143,6 +143,8 @@ def _handle_dataset_data(json_):
             json_.pop("dataset_secret_access_key_list", ()),
             json_.pop("dataset_version_id_list", ()),
             json_.pop("dataset_etag_list", ()),
+            json_.pop("dataset_volume_kind_list", ()),
+            json_.pop("dataset_volume_size_list", ()),
         ]
 
         if not any(datasets):
@@ -165,6 +167,8 @@ def _handle_dataset_data(json_):
                      "aws_secret_access_key": dataset[3],
                      "version_id": dataset[4],
                      "etag": dataset[5],
+                     "volume_kind": dataset[6],
+                     "volume_size": dataset[7],
                      } for dataset in datasets]
 
         json_["datasets"] = datasets

tests/config_files/experiments_create_multi_node.yaml

@@ -14,9 +14,15 @@ datasetAwsSecretAccessKey:
   -
   - some_other_secret
 datasetVersionId:
+  - version1
+  - version2
 datasetEtag:
   - "some etag"
   - "some other etag"
+datasetVolumeKind:
+  - dynamic
+datasetVolumeSize:
+  - 10Gi
 experimentEnv:
   key: val
 experimentType: MPI

tests/config_files/experiments_create_multi_node_ds_objs.yaml

@@ -6,14 +6,16 @@ datasets:
   name: "some dataset name"
   awsAccessKeyId: none
   awsSecretAccessKey: none
-  versionId: version1
   etag: "some etag"
+  volumeKind: "dynamic"
+  volumeSize: "10Gi"
+  versionId: version1
 - uri: "s3://some.other.dataset/uri"
   name: null
   awsAccessKeyId: some_other_key_id
   awsSecretAccessKey: some_other_secret
-  versionId: version2
   etag: "some other etag"
+  versionId: version2
 experimentEnv:
   key: val
 experimentType: MPI