base64-encode commands sent to the APIs so they are not blocked by Cloudflare's filter

(cherry picked from commit 8b6857e)

Author: BartoszCki

Pipfile.lock

@@ -1,7 +1,9 @@
+import copy
+
 import marshmallow
 
-from . import BaseSchema
-from .. import models
+from .base import BaseSchema
+from .. import models, utils
 
 
 class BaseExperimentSchema(BaseSchema):
@@ -22,7 +24,7 @@ class BaseExperimentSchema(BaseSchema):
     id = marshmallow.fields.Str(load_from="handle")
     state = marshmallow.fields.Int()
 
-    def get_instance(self, obj_dict):
+    def get_instance(self, obj_dict, many=False):
         # without popping these marshmallow wouldn't use load_from
         obj_dict.pop("id", None)
         obj_dict.pop("project_id", None)
@@ -31,7 +33,7 @@ def get_instance(self, obj_dict):
         if isinstance(ports, int):
             obj_dict["ports"] = str(ports)
 
-        instance = super(BaseExperimentSchema, self).get_instance(obj_dict)
+        instance = super(BaseExperimentSchema, self).get_instance(obj_dict, many=many)
         return instance
 
 
@@ -46,6 +48,13 @@ class SingleNodeExperimentSchema(BaseExperimentSchema):
     registry_password = marshmallow.fields.Str(dump_to="registryPassword", load_from="registryPassword")
     registry_url = marshmallow.fields.Str(dump_to="registryUrl", load_from="registryUrl")
 
+    @marshmallow.pre_dump
+    def preprocess(self, data, **kwargs):
+        data = copy.copy(data)
+
+        utils.base64_encode_attribute(data, "command")
+        return data
+
 
 class MultiNodeExperimentSchema(BaseExperimentSchema):
     MODEL = models.MultiNodeExperiment
@@ -79,6 +88,14 @@ class MultiNodeExperimentSchema(BaseExperimentSchema):
     parameter_server_registry_url = marshmallow.fields.Str(dump_to="parameterServerRegistryUrl",
                                                            load_from="parameterServerRegistryUrl")
 
+    @marshmallow.pre_dump
+    def preprocess(self, data, **kwargs):
+        data = copy.copy(data)
+
+        utils.base64_encode_attribute(data, "worker_command")
+        utils.base64_encode_attribute(data, "parameter_server_command")
+        return data
+
 
 class MpiMultiNodeExperimentSchema(BaseExperimentSchema):
     MODEL = models.MpiMultiNodeExperiment
@@ -92,10 +109,8 @@ class MpiMultiNodeExperimentSchema(BaseExperimentSchema):
                                               load_from="masterContainer")
     master_machine_type = marshmallow.fields.Str(required=True, dump_to="masterMachineType",
                                                  load_from="masterMachineType")
-    master_command = marshmallow.fields.Str(required=True, dump_to="masterCommand",
-                                            load_from="masterCommand")
-    master_count = marshmallow.fields.Int(required=True, dump_to="masterCount",
-                                          load_from="masterCount")
+    master_command = marshmallow.fields.Str(required=True, dump_to="masterCommand", load_from="masterCommand")
+    master_count = marshmallow.fields.Int(required=True, dump_to="masterCount", load_from="masterCount")
     worker_container_user = marshmallow.fields.Str(dump_to="workerContainerUser", load_from="workerContainerUser")
     worker_registry_username = marshmallow.fields.Str(dump_to="workerRegistryUsername",
                                                       load_from="workerRegistryUsername")
@@ -111,3 +126,11 @@ class MpiMultiNodeExperimentSchema(BaseExperimentSchema):
                                                       load_from="masterRegistryPassword")
     master_registry_url = marshmallow.fields.Str(dump_to="masterRegistryUrl",
                                                  load_from="masterRegistryUrl")
+
+    @marshmallow.pre_dump
+    def preprocess(self, data, **kwargs):
+        data = copy.copy(data)
+
+        utils.base64_encode_attribute(data, "worker_command")
+        utils.base64_encode_attribute(data, "master_command")
+        return data

gradient/api_sdk/serializers/experiment.py

@@ -1,7 +1,9 @@
+import copy
+
 import marshmallow
 
 from .experiment import BaseExperimentSchema
-from .. import models
+from .. import models, utils
 
 
 class HyperparameterSchema(BaseExperimentSchema):
@@ -28,3 +30,11 @@ class HyperparameterSchema(BaseExperimentSchema):
                                                              load_from="hyperparameterServerContainer")
     hyperparameter_server_container_user = marshmallow.fields.Str(dump_to="hyperparameterServerContainerUser",
                                                                   load_from="hyperparameterServerContainerUser")
+
+    @marshmallow.pre_dump
+    def preprocess(self, data, **kwargs):
+        data = copy.copy(data)
+
+        utils.base64_encode_attribute(data, "worker_command")
+        utils.base64_encode_attribute(data, "tuning_command")
+        return data

gradient/api_sdk/serializers/hyperparameter.py

@@ -1,3 +1,4 @@
+import base64
 from collections import OrderedDict
 
 import six
@@ -48,3 +49,20 @@ def print_dict_recursive(input_dict, logger, indent=0, tabulator="  "):
             print_dict_recursive(OrderedDict(val), logger, indent + 1)
         else:
             logger.log("%s%s" % (tabulator * (indent + 1), val))
+
+
+def base64_encode(s):
+    if six.PY3:
+        s = bytes(s, encoding="utf8")
+
+    encoded_str = base64.b64encode(s)
+
+    if six.PY3:  # Python3's base64.b64encode returns a bytes instance so it should be converted back to unicode
+        encoded_str = encoded_str.decode("utf-8")
+
+    return encoded_str
+
+
+def base64_encode_attribute(data, name):
+    encoded_value = base64_encode(getattr(data, name))
+    setattr(data, name, encoded_value)

gradient/api_sdk/utils.py

@@ -56,7 +56,7 @@ class TestExperimentsCreateSingleNode(object):
         "projectHandle": u"testHandle",
         "container": u"testContainer",
         "machineType": u"testType",
-        "command": u"testCommand",
+        "command": u"dGVzdENvbW1hbmQ=",
         "experimentTypeId": constants.ExperimentType.SINGLE_NODE,
         "workspaceUrl": u"some-workspace",
     }
@@ -73,7 +73,7 @@ class TestExperimentsCreateSingleNode(object):
         "projectHandle": u"testHandle",
         "container": u"testContainer",
         "machineType": u"testType",
-        "command": u"testCommand",
+        "command": u"dGVzdENvbW1hbmQ=",
         "containerUser": u"conUser",
         "registryUsername": u"userName",
         "registryPassword": u"passwd",
@@ -333,11 +333,11 @@ class TestExperimentsCreateMultiNode(object):
         u"experimentTypeId": 2,
         u"workerContainer": u"wcon",
         u"workerMachineType": u"mty",
-        u"workerCommand": u"wcom",
+        u"workerCommand": u"d2NvbQ==",
         u"workerCount": 2,
         u"parameterServerContainer": u"pscon",
         u"parameterServerMachineType": u"psmtype",
-        u"parameterServerCommand": u"ls",
+        u"parameterServerCommand": u"bHM=",
         u"parameterServerCount": 2,
         u"workerContainerUser": u"usr",
         u"workspaceUrl": u"https://github.com/Paperspace/gradient-cli.git",
@@ -356,11 +356,11 @@ class TestExperimentsCreateMultiNode(object):
         "experimentTypeId": 3,
         "workerContainer": u"wcon",
         "workerMachineType": u"mty",
-        "workerCommand": u"wcom",
+        "workerCommand": u"d2NvbQ==",
         "workerCount": 2,
         "masterContainer": u"pscon",
         "masterMachineType": u"psmtype",
-        "masterCommand": u"ls",
+        "masterCommand": u"bHM=",
         "masterCount": 2,
         "workerContainerUser": u"usr",
         "workerRegistryUsername": u"rusr",

tests/functional/test_experiments.py

@@ -24,9 +24,9 @@ class TestCreateHyperparameters(object):
         "workerContainer": "some_container",
         "workerMachineType": "k80",
         "name": "some_name",
-        "tuningCommand": "some command",
+        "tuningCommand": "c29tZSBjb21tYW5k",
         "workerCount": 1,
-        "workerCommand": "some worker command",
+        "workerCommand": "c29tZSB3b3JrZXIgY29tbWFuZA==",
         "experimentTypeId": constants.ExperimentType.HYPERPARAMETER_TUNING,
         "projectHandle": "pr4yxj956",
     }
@@ -66,9 +66,9 @@ class TestCreateHyperparameters(object):
         "workerContainer": "some_worker_container",
         "workerMachineType": "k80",
         "name": "some_name",
-        "tuningCommand": "some command",
+        "tuningCommand": "c29tZSBjb21tYW5k",
         "workerCount": 666,
-        "workerCommand": "some worker command",
+        "workerCommand": "c29tZSB3b3JrZXIgY29tbWFuZA==",
         "workerRegistryUsername": "some_registry_username",
         "workerRegistryPassword": "some_registry_password",
         "workerContainerUser": "some_worker_container_user",
@@ -273,9 +273,9 @@ class TestCreateAndStartHyperparameters(object):
         "workerContainer": "some_container",
         "workerMachineType": "k80",
         "name": "some_name",
-        "tuningCommand": "some command",
+        "tuningCommand": "c29tZSBjb21tYW5k",
         "workerCount": 1,
-        "workerCommand": "some worker command",
+        "workerCommand": "c29tZSB3b3JrZXIgY29tbWFuZA==",
         "projectHandle": "pr4yxj956",
         "experimentTypeId": constants.ExperimentType.HYPERPARAMETER_TUNING,
     }
@@ -315,9 +315,9 @@ class TestCreateAndStartHyperparameters(object):
         "workerContainer": "some_worker_container",
         "workerMachineType": "k80",
         "name": "some_name",
-        "tuningCommand": "some command",
+        "tuningCommand": "c29tZSBjb21tYW5k",
         "workerCount": 666,
-        "workerCommand": "some worker command",
+        "workerCommand": "c29tZSB3b3JrZXIgY29tbWFuZA==",
         "workerRegistryUsername": "some_registry_username",
         "workerRegistryPassword": "some_registry_password",
         "workerContainerUser": "some_worker_container_user",