Add optional --apiKey to experiments commands and refactor creating client.API instance

Author: BartoszCki

paperspace/cli.py

@@ -3,7 +3,7 @@
 
 import click
 
-from paperspace import constants
+from paperspace import constants, client, config
 from paperspace.commands import experiments as experiments_commands
 
 
@@ -113,6 +113,10 @@ def common_experiments_create_options(f):
             "--modelPath",
             "modelPath",
         ),
+        click.option(
+            "--apiKey",
+            "api_key",
+        )
     ]
     return functools.reduce(lambda x, opt: opt(x), reversed(options), f)
 
@@ -229,60 +233,100 @@ def common_experiments_create_single_node_options(f):
 @create.command(name="multinode")
 @common_experiments_create_options
 @common_experiment_create_multi_node_options
-def create_multi_node(**kwargs):
+def create_multi_node(api_key, **kwargs):
     del_if_value_is_none(kwargs)
-    experiments_commands.create_experiment(kwargs)
+    experiments_api = client.API(config.CONFIG_EXPERIMENTS_HOST)
+    if api_key:
+        experiments_api.api_key = api_key
+    experiments_commands.create_experiment(kwargs, api=experiments_api)
 
 
 @create.command(name="singlenode")
 @common_experiments_create_options
 @common_experiments_create_single_node_options
-def create_single_node(**kwargs):
+def create_single_node(api_key, **kwargs):
     kwargs["experimentTypeId"] = constants.ExperimentType.SINGLE_NODE
     del_if_value_is_none(kwargs)
-    experiments_commands.create_experiment(kwargs)
+    experiments_api = client.API(config.CONFIG_EXPERIMENTS_HOST)
+    if api_key:
+        experiments_api.api_key = api_key
+    experiments_commands.create_experiment(kwargs, api=experiments_api)
 
 
 @create_and_start.command(name="multinode")
 @common_experiments_create_options
 @common_experiment_create_multi_node_options
-def create_and_start_multi_node(**kwargs):
+def create_and_start_multi_node(api_key, **kwargs):
     del_if_value_is_none(kwargs)
-    experiments_commands.create_and_start_experiment(kwargs)
+    experiments_api = client.API(config.CONFIG_EXPERIMENTS_HOST)
+    if api_key:
+        experiments_api.api_key = api_key
+    experiments_commands.create_and_start_experiment(kwargs, api=experiments_api)
 
 
 @create_and_start.command(name="singlenode")
 @common_experiments_create_options
 @common_experiments_create_single_node_options
-def create_and_start_single_node(**kwargs):
+def create_and_start_single_node(api_key, **kwargs):
     kwargs["experimentTypeId"] = constants.ExperimentType.SINGLE_NODE
     del_if_value_is_none(kwargs)
-    experiments_commands.create_and_start_experiment(kwargs)
+    experiments_api = client.API(config.CONFIG_EXPERIMENTS_HOST)
+    if api_key:
+        experiments_api.api_key = api_key
+    experiments_commands.create_and_start_experiment(kwargs, api=experiments_api)
 
 
 @experiments.command()
 @click.argument("experiment-handle")
-def start(experiment_handle):
-    experiments_commands.start_experiment(experiment_handle)
+@click.option(
+    "--apiKey",
+    "api_key",
+)
+def start(experiment_handle, api_key):
+    experiments_api = client.API(config.CONFIG_EXPERIMENTS_HOST)
+    if api_key:
+        experiments_api.api_key = api_key
+    experiments_commands.start_experiment(experiment_handle, api=experiments_api)
 
 
 @experiments.command()
 @click.argument("experiment-handle")
-def stop(experiment_handle):
-    experiments_commands.stop_experiment(experiment_handle)
+@click.option(
+    "--apiKey",
+    "api_key",
+)
+def stop(experiment_handle, api_key):
+    experiments_api = client.API(config.CONFIG_EXPERIMENTS_HOST)
+    if api_key:
+        experiments_api.api_key = api_key
+    experiments_commands.stop_experiment(experiment_handle, api=experiments_api)
 
 
 @experiments.command("list")
 @click.option("--projectHandle", "-p", "project_handles", multiple=True)
-def list_experiments(project_handles):
-    command = experiments_commands.ListExperimentsCommand()
+@click.option(
+    "--apiKey",
+    "api_key",
+)
+def list_experiments(project_handles, api_key):
+    experiments_api = client.API(config.CONFIG_EXPERIMENTS_HOST)
+    if api_key:
+        experiments_api.api_key = api_key
+    command = experiments_commands.ListExperimentsCommand(api=experiments_api)
     command.execute(project_handles)
 
 
 @experiments.command("details")
 @click.argument("experiment-handle")
-def get_experiment_details(experiment_handle):
-    experiments_commands.get_experiment_details(experiment_handle)
+@click.option(
+    "--apiKey",
+    "api_key",
+)
+def get_experiment_details(experiment_handle, api_key):
+    experiments_api = client.API(config.CONFIG_EXPERIMENTS_HOST)
+    if api_key:
+        experiments_api.api_key = api_key
+    experiments_commands.get_experiment_details(experiment_handle, api=experiments_api)
 
 # TODO: delete experiment - not implemented in the api
 # TODO: modify experiment - not implemented in the api

paperspace/client.py

@@ -1,12 +1,17 @@
 import requests
 
-from paperspace import logger
+from paperspace import logger, config, version
+
+default_headers = {"X-API-Key": config.PAPERSPACE_API_KEY,
+                   "ps_client_name": "paperspace-python",
+                   "ps_client_version": version.version}
 
 
 class API(object):
     def __init__(self, api_url, headers=None):
         self.api_url = api_url
-        self.headers = headers or {}
+        headers = headers or default_headers
+        self.headers = headers.copy()
 
     def get_path(self, url):
         api_url = self.api_url if not self.api_url.endswith("/") else self.api_url[:-1]
@@ -36,3 +41,11 @@ def get(self, url, params=None):
         logger.debug("Response status code: {}".format(response.status_code))
         logger.debug("Response content: {}".format(response.content))
         return response
+
+    @property
+    def api_key(self):
+        return self.headers.get("X-API-Key")
+
+    @api_key.setter
+    def api_key(self, value):
+        self.headers["X-API-Key"] = value

paperspace/commands/experiments.py

@@ -2,14 +2,11 @@
 
 import terminaltables
 
-from paperspace import config, version, client, logger, constants
+from paperspace import logger, constants, client, config
 from paperspace.logger import log_response
 from paperspace.utils import get_terminal_lines
 
-default_headers = {"X-API-Key": config.PAPERSPACE_API_KEY,
-                   "ps_client_name": "paperspace-python",
-                   "ps_client_version": version.version}
-experiments_api = client.API(config.CONFIG_EXPERIMENTS_HOST, headers=default_headers)
+experiments_api = client.API(config.CONFIG_EXPERIMENTS_HOST, headers=client.default_headers)
 
 
 def _log_create_experiment(response, success_msg_template, error_msg, logger_=logger):
@@ -183,4 +180,4 @@ def get_experiment_details(experiment_handle, api=experiments_api):
             logger.log("Error parsing response data")
             logger.debug(e)
 
-    log_response(response, details, "Unknown error while retrieving details of the experiment")
+    log_response(response, details, "Unknown error while retrieving details of the experiment")

tests/test_click_commands.py

@@ -4,16 +4,20 @@
 from paperspace import cli, constants
 
 
+@mock.patch("paperspace.cli.client.API")
 @mock.patch("paperspace.cli.experiments_commands")
-def test_should_execute_create_experiment_command_when_cli_singlenode_command_was_executed(commands_patched):
+def test_should_execute_create_experiment_command_when_cli_singlenode_command_was_executed(commands_patched,
+                                                                                           api_patched):
+    api_patched.return_value = mock.MagicMock()
     runner = CliRunner()
     command = "experiments create singlenode " \
               "--name exp1 " \
               "--projectHandle testHandle " \
               "--container testContainer " \
               "--machineType testType " \
               "--command testCommand " \
-              "--workspaceUrl wUrl"
+              "--workspaceUrl wUrl " \
+              "--apiKey some_key"
     expected_kwargs = {"name": u"exp1",
                        "projectHandle": u"testHandle",
                        "container": u"testContainer",
@@ -26,11 +30,14 @@ def test_should_execute_create_experiment_command_when_cli_singlenode_command_wa
     result = runner.invoke(cli.cli, command.split())
 
     assert result.exit_code == 0
-    commands_patched.create_experiment.assert_called_once_with(expected_kwargs)
+    commands_patched.create_experiment.assert_called_once_with(expected_kwargs, api=api_patched())
 
 
+@mock.patch("paperspace.cli.client.API")
 @mock.patch("paperspace.cli.experiments_commands")
-def test_should_execute_create_experiment_command_when_cli_multinode_mpi_command_was_executed(commands_patched):
+def test_should_execute_create_experiment_command_when_cli_multinode_mpi_command_was_executed(commands_patched,
+                                                                                              api_patched):
+    api_patched.return_value = mock.MagicMock()
     runner = CliRunner()
     command = "experiments create multinode " \
               "--name exp1 " \
@@ -44,7 +51,8 @@ def test_should_execute_create_experiment_command_when_cli_multinode_mpi_command
               "--parameterServerMachineType testParameterServerMachineType " \
               "--parameterServerCommand testParameterServerCommand " \
               "--parameterServerCount 3 " \
-              "--workspaceUrl wUrl"
+              "--workspaceUrl wUrl " \
+              "--apiKey some_key"
     expected_kwargs = {"name": u"exp1",
                        "projectHandle": u"testHandle",
                        "experimentTypeId": constants.ExperimentType.MPI_MULTI_NODE,
@@ -62,11 +70,14 @@ def test_should_execute_create_experiment_command_when_cli_multinode_mpi_command
     result = runner.invoke(cli.cli, command.split())
 
     assert result.exit_code == 0
-    commands_patched.create_experiment.assert_called_once_with(expected_kwargs)
+    commands_patched.create_experiment.assert_called_once_with(expected_kwargs, api=api_patched())
 
 
+@mock.patch("paperspace.cli.client.API")
 @mock.patch("paperspace.cli.experiments_commands")
-def test_should_execute_create_experiment_command_when_cli_multinode_grpc_command_was_executed(commands_patched):
+def test_should_execute_create_experiment_command_when_cli_multinode_grpc_command_was_executed(commands_patched,
+                                                                                               api_patched):
+    api_patched.return_value = mock.MagicMock()
     runner = CliRunner()
     command = "experiments create multinode " \
               "--name exp1 " \
@@ -98,20 +109,23 @@ def test_should_execute_create_experiment_command_when_cli_multinode_grpc_comman
     result = runner.invoke(cli.cli, command.split())
 
     assert result.exit_code == 0
-    commands_patched.create_experiment.assert_called_once_with(expected_kwargs)
+    commands_patched.create_experiment.assert_called_once_with(expected_kwargs, api=api_patched())
 
 
+@mock.patch("paperspace.cli.client.API")
 @mock.patch("paperspace.cli.experiments_commands")
 def test_should_execute_create_experiment_command_when_cli_create_and_start_singlenode_command_was_executed(
-        commands_patched):
+        commands_patched, api_patched):
+    api_patched.return_value = mock.MagicMock()
     runner = CliRunner()
     command = "experiments createAndStart singlenode " \
               "--name exp1 " \
               "--projectHandle testHandle " \
               "--container testContainer " \
               "--machineType testType " \
               "--command testCommand " \
-              "--workspaceUrl wUrl"
+              "--workspaceUrl wUrl " \
+              "--apiKey some_key"
     expected_kwargs = {"name": u"exp1",
                        "projectHandle": u"testHandle",
                        "container": u"testContainer",
@@ -124,12 +138,14 @@ def test_should_execute_create_experiment_command_when_cli_create_and_start_sing
     result = runner.invoke(cli.cli, command.split())
 
     assert result.exit_code == 0
-    commands_patched.create_and_start_experiment.assert_called_once_with(expected_kwargs)
+    commands_patched.create_and_start_experiment.assert_called_once_with(expected_kwargs, api=api_patched())
 
 
+@mock.patch("paperspace.cli.client.API")
 @mock.patch("paperspace.cli.experiments_commands")
 def test_should_execute_create_experiment_command_when_cli_create_and_start_multinode_mpi_command_was_executed(
-        commands_patched):
+        commands_patched, api_patched):
+    api_patched.return_value = mock.MagicMock()
     runner = CliRunner()
     command = "experiments createAndStart multinode " \
               "--name exp1 " \
@@ -161,4 +177,4 @@ def test_should_execute_create_experiment_command_when_cli_create_and_start_mult
     result = runner.invoke(cli.cli, command.split())
 
     assert result.exit_code == 0
-    commands_patched.create_and_start_experiment.assert_called_once_with(expected_kwargs)
+    commands_patched.create_and_start_experiment.assert_called_once_with(expected_kwargs, api=api_patched())