Merge pull request #93 from Paperspace/PS-9792

Ps 9792

Author: mkulaczkowski

paperspace/cli/jobs.py

@@ -1,10 +1,9 @@
 import click
 
 from paperspace import client, config
-from paperspace.cli import common
-from paperspace.cli.cli_types import json_string
-from paperspace.cli.common import del_if_value_is_none, ClickGroup
 from paperspace.cli.cli import cli
+from paperspace.cli.cli_types import json_string
+from paperspace.cli.common import api_key_option, del_if_value_is_none, ClickGroup
 from paperspace.commands import jobs as jobs_commands
 
 
@@ -20,7 +19,7 @@ def jobs_group():
     required=True,
     help="Delete job with given ID",
 )
-@common.api_key_option
+@api_key_option
 def delete_job(job_id, api_key=None):
     jobs_api = client.API(config.CONFIG_HOST, api_key=api_key)
     command = jobs_commands.DeleteJobCommand(api=jobs_api)
@@ -34,7 +33,7 @@ def delete_job(job_id, api_key=None):
     required=True,
     help="Stop job with given ID",
 )
-@common.api_key_option
+@api_key_option
 def stop_job(job_id, api_key=None):
     jobs_api = client.API(config.CONFIG_HOST, api_key=api_key)
     command = jobs_commands.StopJobCommand(api=jobs_api)
@@ -57,9 +56,9 @@ def stop_job(job_id, api_key=None):
     "experimentId",
     help="Use to filter jobs by experiment ID",
 )
-@common.api_key_option
+@api_key_option
 def list_jobs(api_key, **filters):
-    common.del_if_value_is_none(filters)
+    del_if_value_is_none(filters)
     jobs_api = client.API(config.CONFIG_HOST, api_key=api_key)
     command = jobs_commands.ListJobsCommand(api=jobs_api)
     command.execute(filters=filters)
@@ -87,7 +86,7 @@ def list_jobs(api_key, **filters):
 @click.option("--relDockerfilePath", "relDockerfilePath", help="Relative path to Dockerfile")
 @click.option("--registryUsername", "registryUsername", help="Docker registry username")
 @click.option("--registryPassword", "registryPassword", help="Docker registry password")
-@common.api_key_option
+@api_key_option
 def create_job(api_key, **kwargs):
     del_if_value_is_none(kwargs)
     jobs_api = client.API(config.CONFIG_HOST, api_key=api_key)
@@ -101,8 +100,44 @@ def create_job(api_key, **kwargs):
     "job_id",
     required=True
 )
-@common.api_key_option
+@api_key_option
 def list_logs(job_id, api_key=None):
     logs_api = client.API(config.CONFIG_LOG_HOST, api_key=api_key)
     command = jobs_commands.JobLogsCommand(api=logs_api)
     command.execute(job_id)
+
+
+@jobs_group.group("artifacts", help="Manage jobs' artifacts", cls=ClickGroup)
+def artifacts():
+    pass
+
+
+@artifacts.command("destroy", help="Destroy job's artifacts")
+@click.argument("job_id")
+@click.option("--files", "files")
+@api_key_option
+def destroy_artifacts(job_id, api_key=None, files=None):
+    jobs_api = client.API(config.CONFIG_HOST, api_key=api_key)
+    command = jobs_commands.ArtifactsDestroyCommand(api=jobs_api)
+    command.execute(job_id, files=files)
+
+
+@artifacts.command("get", help="Get job's artifacts")
+@click.argument("job_id")
+@api_key_option
+def get_artifacts(job_id, api_key=None):
+    jobs_api = client.API(config.CONFIG_HOST, api_key=api_key)
+    command = jobs_commands.ArtifactsGetCommand(api=jobs_api)
+    command.execute(job_id)
+
+
+@artifacts.command("list", help="List job's artifacts")
+@click.argument("job_id")
+@click.option("--size", "-s", "size", help="Show file size", is_flag=True)
+@click.option("--links", "-l", "links", help="Show file URL", is_flag=True)
+@click.option("--files", "files", help="Get only given file (use at the end * as a wildcard)")
+@api_key_option
+def list_artifacts(job_id, size, links, files, api_key=None):
+    jobs_api = client.API(config.CONFIG_HOST, api_key=api_key)
+    command = jobs_commands.ArtifactsListCommand(api=jobs_api)
+    command.execute(job_id=job_id, size=size, links=links, files=files)

paperspace/commands/__init__.py

@@ -0,0 +1,17 @@
+from collections import OrderedDict
+
+from paperspace import logger
+
+
+class CommandBase(object):
+    def __init__(self, api=None, logger_=logger):
+        self.api = api
+        self.logger = logger_
+
+    def _print_dict_recursive(self, input_dict, indent=0, tabulator="  "):
+        for key, val in input_dict.items():
+            self.logger.log("%s%s:" % (tabulator * indent, key))
+            if type(val) is dict:
+                self._print_dict_recursive(OrderedDict(val), indent + 1)
+            else:
+                self.logger.log("%s%s" % (tabulator * (indent + 1), val))

paperspace/commands/jobs.py

@@ -4,9 +4,10 @@
 from click import style
 
 from paperspace import config, client
+from paperspace.commands import common
+from paperspace.exceptions import BadResponseError
 from paperspace.utils import get_terminal_lines
 from paperspace.workspace import S3WorkspaceHandler
-from paperspace.commands import common
 
 
 class JobsCommandBase(common.CommandBase):
@@ -140,3 +141,76 @@ def execute(self, json_):
         self._log_message(response,
                           "Job created",
                           "Unknown error while creating job")
+
+
+class ArtifactsDestroyCommand(JobsCommandBase):
+    def execute(self, job_id, files=None):
+        url = '/jobs/{}/artifactsDestroy'.format(job_id)
+        params = None
+        if files:
+            params = {'files': files}
+        response = self.api.post(url, params=params)
+        self._log_message(response, "Artifacts destroyed", "Unknown error while destroying artifacts")
+
+
+class ArtifactsGetCommand(JobsCommandBase):
+    def execute(self, job_id):
+        url = '/jobs/artifactsGet'
+        response = self.api.get(url, params={'jobId': job_id})
+
+        self._log_artifacts(response)
+
+    def _log_artifacts(self, response):
+        try:
+            artifacts_json = response.json()
+            if response.ok:
+                self._print_dict_recursive(artifacts_json)
+            else:
+                raise BadResponseError(
+                    '{}: {}'.format(artifacts_json['error']['status'], artifacts_json['error']['message']))
+        except (ValueError, KeyError, BadResponseError) as e:
+            self.logger.error("Error occurred while getting artifacts: {}".format(str(e)))
+
+
+class ArtifactsListCommand(common.ListCommand):
+    kwargs = {}
+
+    def execute(self, **kwargs):
+        self.kwargs = kwargs
+        return super(ArtifactsListCommand, self).execute(**kwargs)
+
+    @property
+    def request_url(self):
+        return '/jobs/artifactsList'
+
+    def _get_request_params(self, kwargs):
+        params = {'jobId': kwargs['job_id']}
+
+        files = kwargs.get('files')
+        if files:
+            params['files'] = files
+        size = kwargs.get('size', False)
+        if size:
+            params['size'] = size
+        links = kwargs.get('links', False)
+        if links:
+            params['links'] = links
+
+        return params
+
+    def _get_table_data(self, artifacts):
+        columns = ['Files']
+        if self.kwargs.get('size'):
+            columns.append('Size (in bytes)')
+        if self.kwargs.get('links'):
+            columns.append('URL')
+
+        data = [tuple(columns)]
+        for artifact in artifacts:
+            row = [artifact.get('file')]
+            if 'size' in artifact.keys():
+                row.append(artifact['size'])
+            if 'url' in artifact.keys():
+                row.append(artifact['url'])
+            data.append(tuple(row))
+        return data

tests/functional/test_jobs.py

@@ -1,15 +1,23 @@
 import mock
+import pytest
 from click.testing import CliRunner
 
-import paperspace
 from paperspace.cli import cli
 from paperspace.client import default_headers
 from tests import example_responses, MockResponse
 
 
-class TestListJobs(object):
-    URL = "https://api.paperspace.io/jobs/getJobs/"
+class TestJobs(object):
+    EXPECTED_STDOUT_WITH_WRONG_API_TOKEN = "Invalid API token\n"
+    RESPONSE_JSON_WITH_WRONG_API_TOKEN = {"status": 400, "message": "Invalid API token"}
+
     EXPECTED_HEADERS = default_headers.copy()
+    EXPECTED_HEADERS_WITH_CHANGED_API_KEY = default_headers.copy()
+    EXPECTED_HEADERS_WITH_CHANGED_API_KEY["X-API-Key"] = "some_key"
+
+
+class TestListJobs(TestJobs):
+    URL = "https://api.paperspace.io/jobs/getJobs/"
     BASIC_COMMAND = ["jobs", "list"]
     EXPECTED_RESPONSE_JSON = example_responses.LIST_JOBS_RESPONSE_JSON
     EXPECTED_STDOUT = """+----------------+---------------------------+-------------------+----------------+--------------+--------------------------+
@@ -30,11 +38,6 @@ class TestListJobs(object):
 """
 
     BASIC_COMMAND_WITH_API_KEY = ["jobs", "list", "--apiKey", "some_key"]
-    EXPECTED_HEADERS_WITH_CHANGED_API_KEY = paperspace.client.default_headers.copy()
-    EXPECTED_HEADERS_WITH_CHANGED_API_KEY["X-API-Key"] = "some_key"
-
-    RESPONSE_JSON_WITH_WRONG_API_TOKEN = {"status": 400, "message": "Invalid API token"}
-    EXPECTED_STDOUT_WITH_WRONG_API_TOKEN = "Invalid API token\n"
 
     RESPONSE_JSON_WHEN_NO_JOBS_WERE_FOUND = []
     EXPECTED_STDOUT_WHEN_NO_JOBS_WERE_FOUND = "No data found\n"
@@ -169,11 +172,8 @@ def test_should_print_proper_message_when_jobs_list_was_used_with_mutually_exclu
         assert result.exit_code == 0
 
 
-class TestJobLogs(object):
+class TestJobLogs(TestJobs):
     URL = "https://logs.paperspace.io/jobs/logs?jobId=some_job_id&line=0"
-    EXPECTED_HEADERS = default_headers.copy()
-    EXPECTED_HEADERS_WITH_CHANGED_API_KEY = default_headers.copy()
-    EXPECTED_HEADERS_WITH_CHANGED_API_KEY["X-API-Key"] = "some_key"
 
     RESPONSE_JSON_WITH_WRONG_API_TOKEN = {"status": 400, "message": "Invalid API token"}
     EXPECTED_RESPONSE_JSON = example_responses.LIST_OF_LOGS_FOR_JOB
@@ -257,3 +257,86 @@ def test_should_print_error_message_when_error_status_code_received_but_no_conte
                                        params=None)
         assert result.output == "Error while parsing response data: No JSON\n"
         assert result.exit_code == 0
+
+
+class TestJobArtifactsCommands(TestJobs):
+    runner = CliRunner()
+    URL = "https://api.paperspace.io"
+
+    @mock.patch("paperspace.client.requests.post")
+    def test_should_send_valid_post_request_when_destroying_artifacts_with_files_specified(self, post_patched):
+        post_patched.return_value = MockResponse(status_code=200)
+        job_id = "some_job_id"
+        file_names = "some_file_names"
+        result = self.runner.invoke(cli.cli, ["jobs", "artifacts", "destroy", job_id, "--files", file_names, "--apiKey",
+                                              "some_key"])
+
+        post_patched.assert_called_with("{}/jobs/{}/artifactsDestroy".format(self.URL, job_id),
+                                        files=None,
+                                        headers=self.EXPECTED_HEADERS_WITH_CHANGED_API_KEY,
+                                        json=None,
+                                        params={"files": file_names})
+        assert result.exit_code == 0
+
+    @mock.patch("paperspace.client.requests.post")
+    def test_should_send_valid_post_request_when_destroying_artifacts_without_files_specified(self, post_patched):
+        post_patched.return_value = MockResponse(status_code=200)
+        job_id = "some_job_id"
+        result = self.runner.invoke(cli.cli, ["jobs", "artifacts", "destroy", job_id, "--apiKey", "some_key"])
+
+        post_patched.assert_called_with("{}/jobs/{}/artifactsDestroy".format(self.URL, job_id),
+                                        files=None,
+                                        headers=self.EXPECTED_HEADERS_WITH_CHANGED_API_KEY,
+                                        json=None,
+                                        params=None)
+        assert result.exit_code == 0
+
+    @mock.patch("paperspace.client.requests.get")
+    def test_should_send_send_valid_get_request_and_receive_json_response(self, get_patched):
+        get_patched.return_value = MockResponse(status_code=200)
+        job_id = "some_job_id"
+        result = self.runner.invoke(cli.cli, ["jobs", "artifacts", "get", job_id, "--apiKey", "some_key"])
+
+        get_patched.assert_called_with("{}/jobs/artifactsGet".format(self.URL),
+                                       headers=self.EXPECTED_HEADERS_WITH_CHANGED_API_KEY,
+                                       json=None,
+                                       params={"jobId": job_id})
+        assert result.exit_code == 0
+
+    @mock.patch("paperspace.client.requests.get")
+    def test_should_send_valid_get_request_with_all_parameters_for_a_list_of_artifacts(self, get_patched):
+        get_patched.return_value = MockResponse(status_code=200)
+        job_id = "some_job_id"
+        result = self.runner.invoke(cli.cli,
+                                    ["jobs", "artifacts", "list", job_id, "--apiKey", "some_key", "--size", "--links",
+                                     "--files", "foo"])
+
+        get_patched.assert_called_with("{}/jobs/artifactsList".format(self.URL),
+                                       headers=self.EXPECTED_HEADERS_WITH_CHANGED_API_KEY,
+                                       json=None,
+                                       params={"jobId": job_id,
+                                               "size": True,
+                                               "links": True,
+                                               "files": "foo"})
+        assert result.exit_code == 0
+
+    @mock.patch("paperspace.client.requests.get")
+    @pytest.mark.parametrize('option,param', [("--size", "size"),
+                                              ("-s", "size"),
+                                              ("--links", "links"),
+                                              ("-l", "links")])
+    def test_should_send_valid_get_request_with_valid_param_for_a_list_of_artifacts_for_both_formats_of_param(self,
+                                                                                                              get_patched,
+                                                                                                              option,
+                                                                                                              param):
+        get_patched.return_value = MockResponse(status_code=200)
+        job_id = "some_job_id"
+        result = self.runner.invoke(cli.cli,
+                                    ["jobs", "artifacts", "list", job_id, "--apiKey", "some_key"] + [option])
+
+        get_patched.assert_called_with("{}/jobs/artifactsList".format(self.URL),
+                                       headers=self.EXPECTED_HEADERS_WITH_CHANGED_API_KEY,
+                                       json=None,
+                                       params={"jobId": job_id,
+                                               param: True})
+        assert result.exit_code == 0

tests/unit/test_base.py

@@ -0,0 +1,39 @@
+from collections import OrderedDict
+
+import mock
+import pytest
+
+from paperspace.commands import CommandBase
+
+output_response = ""
+
+
+class TestBaseClass(object):
+    def test_json_print(self):
+        global output_response
+        output_response = ""
+
+        def log_to_var(message):
+            global output_response
+            output_response = "{}{}\n".format(output_response, message)
+
+        logger_ = mock.MagicMock()
+        logger_.log = log_to_var
+
+        input_dict = {
+            "foo": {
+                'bar': {
+                    "baz": "faz"
+                }
+            }
+        }
+        expected_string = """foo:
+  bar:
+    baz:
+      faz
+"""
+
+        command = CommandBase(logger_=logger_)
+        command._print_dict_recursive(OrderedDict(input_dict))
+
+        assert output_response == expected_string