Add support for listJobArtifactsV2 endpoint to support pagination PS-15028 (#326)

Author: paperspace-philip

gradient/api_sdk/clients/job_client.py

@@ -342,7 +342,7 @@ def artifacts_get(self, job_id):
         data = repository.get(jobId=job_id)
         return data
 
-    def artifacts_list(self, job_id, files=None, size=False, links=True):
+    def artifacts_list(self, job_id, files=None, size=False, links=True, start_after=None):
         """
         Method to retrieve all artifacts files.
 
@@ -354,20 +354,21 @@ def artifacts_list(self, job_id, files=None, size=False, links=True):
                 job_id='your_job_id_here',
                 files='your_files,here',
                 size=False,
-                links=True
+                links=True,
+                start_after='key',
             )
 
         :param str job_id: to limit artifact from this job.
         :param str files: to limit result only to file names provided. You can use wildcard option ``*``.
         :param bool size: flag to show file size. Default value is set to False.
         :param bool links: flag to show file url. Default value is set to True.
+        :params str start_after: key to list after
 
         :returns: list of files with description if specified from job artifacts.
-        :rtype: list[Artifact]
+        :rtype: Pagination
         """
         repository = self.build_repository(ListJobArtifacts)
-        artifacts = repository.list(jobId=job_id, files=files, links=links, size=size)
-        return artifacts
+        return repository.list(jobId=job_id, files=files, links=links, size=size, start_after=start_after)
 
     def get_metrics(self, job_id, start=None, end=None, interval="30s", built_in_metrics=None):
         """Get job metrics

gradient/api_sdk/models/__init__.py

@@ -9,6 +9,7 @@
 from .machine import Machine, MachineEvent, MachineUtilization
 from .model import Model, ModelFile
 from .notebook import Notebook, NotebookStart
+from .pagination import Pagination
 from .project import Project
 from .secret import Secret
 from .tag import Tag

gradient/api_sdk/models/pagination.py

@@ -0,0 +1,15 @@
+from datetime import datetime
+
+import attr
+
+
+@attr.s
+class Pagination(object):
+    """
+    Pagination class
+
+    :param list[Any] data:
+    :param str start_after:
+    """
+    data = attr.ib(type=list, default=None)
+    start_after = attr.ib(type=str, default=None)

gradient/api_sdk/repositories/jobs.py

@@ -126,12 +126,11 @@ def _parse_object(self, instance_dict, **kwargs):
 
 class ListJobArtifacts(GetBaseJobApiUrlMixin, ListResources):
     def _parse_objects(self, data, **kwargs):
-        serializer = serializers.ArtifactSchema()
-        files = serializer.get_instance(data, many=True)
-        return files
+        serializer = serializers.utils.paginate_schema(serializers.ArtifactSchema)
+        return serializer.get_instance(data)
 
     def get_request_url(self, **kwargs):
-        return "/jobs/artifactsList"
+        return "/jobs/artifactsListV2"
 
     def _get_request_params(self, kwargs):
         params = {
@@ -146,6 +145,9 @@ def _get_request_params(self, kwargs):
 
         if kwargs.get("links"):
             params["links"] = kwargs.get("links")
+        
+        if kwargs.get("start_after"):
+            params["startAfter"] = kwargs.get("start_after")
 
         return params
 

gradient/api_sdk/s3_downloader.py

@@ -107,7 +107,18 @@ class JobArtifactsDownloader(ResourceDownloader):
     CLIENT_CLASS = JobsClient
 
     def _get_files_list(self, job_id):
-        files = self.client.artifacts_list(job_id)
+        start_after = None
+        files = []
+        while True:
+            pagination_response = self.client.artifacts_list(job_id, start_after=start_after)
+
+            if pagination_response.data:
+                files.extend(pagination_response.data)
+            start_after = pagination_response.start_after
+
+            if start_after is None:
+                break
+
         files = tuple((f.file, f.url) for f in files)
         return files
 

gradient/api_sdk/serializers/utils.py

@@ -1,7 +1,13 @@
+import marshmallow
+
+from .base import BaseSchema
+
+from .. import models
 from ..sdk_exceptions import GradientSdkError
 from ..serializers import SingleNodeExperimentSchema, MultiNodeExperimentSchema, HyperparameterSchema, \
     MpiMultiNodeExperimentSchema
 
+
 EXPERIMENT_ID_TO_EXPERIMENT_SERIALIZER_MAPPING = {
     1: SingleNodeExperimentSchema,
     2: MultiNodeExperimentSchema,
@@ -23,3 +29,12 @@ def get_serializer_for_experiment(experiment_dict):
         raise GradientSdkError("No experiment type with ID: {}".format(str(e)))
 
     return serializer
+
+
+def paginate_schema(schema):
+    class PaginationSchema(BaseSchema):
+        MODEL = models.Pagination
+        data = marshmallow.fields.Nested(schema, many=True, dump_only=True)
+        start_after = marshmallow.fields.String(dump_to='startAfter', load_from='startAfter')
+
+    return PaginationSchema()

gradient/commands/jobs.py

@@ -193,7 +193,17 @@ class ArtifactsListCommand(BaseJobCommand):
     def execute(self, **kwargs):
         with halo.Halo(text=self.WAITING_FOR_RESPONSE_MESSAGE, spinner="dots"):
             try:
-                instances = self.client.artifacts_list(**kwargs)
+                start_after = None
+                instances = []
+                while True:
+                    pagination_response = self.client.artifacts_list(start_after=start_after, **kwargs)
+
+                    if pagination_response.data:
+                        instances.extend(pagination_response.data)
+                    start_after = pagination_response.start_after
+
+                    if start_after is None:
+                        break
             except sdk_exceptions.GradientSdkError as e:
                 raise exceptions.ReceivingDataFailedError(e)
 

tests/example_responses.py

@@ -5720,20 +5720,22 @@
     "message": "success"
 }
 
-LIST_JOB_FILES_RESPONSE_JSON = [
-    {
-        "file": "hello.txt",
-        "url": "https://ps-projects.s3.amazonaws.com/some/path/artifacts/hello.txt?AWSAccessKeyId=some_aws_access_key_id&Expires=713274132&Signature=7CT5k6buEmZe5k5E7g6BXMs2xV4%3D&response-content-disposition=attachment%3Bfilename%3D%22hello.txt%22&x-amz-security-token=some_amz_security_token"
-    },
-    {
-        "file": "hello2.txt",
-        "url": "https://ps-projects.s3.amazonaws.com/some/path/artifacts/hello2.txt?AWSAccessKeyId=some_aws_access_key_id&Expires=713274132&Signature=L1lI47cNyiROzdYkf%2FF3Cm3165E%3D&response-content-disposition=attachment%3Bfilename%3D%22hello2.txt%22&x-amz-security-token=some_amz_security_token"
-    },
-    {
-        "file": "keton/elo.txt",
-        "url": "https://ps-projects.s3.amazonaws.com/some/path/artifacts/keton/elo.txt?AWSAccessKeyId=some_aws_access_key_id&Expires=713274132&Signature=tHriojGx03S%2FKkVGQGVI5CQRFTo%3D&response-content-disposition=attachment%3Bfilename%3D%22elo.txt%22&x-amz-security-token=some_amz_security_token"
-    }
-]
+LIST_JOB_FILES_RESPONSE_JSON = {
+    'data': [
+        {
+            "file": "hello.txt",
+            "url": "https://ps-projects.s3.amazonaws.com/some/path/artifacts/hello.txt?AWSAccessKeyId=some_aws_access_key_id&Expires=713274132&Signature=7CT5k6buEmZe5k5E7g6BXMs2xV4%3D&response-content-disposition=attachment%3Bfilename%3D%22hello.txt%22&x-amz-security-token=some_amz_security_token"
+        },
+        {
+            "file": "hello2.txt",
+            "url": "https://ps-projects.s3.amazonaws.com/some/path/artifacts/hello2.txt?AWSAccessKeyId=some_aws_access_key_id&Expires=713274132&Signature=L1lI47cNyiROzdYkf%2FF3Cm3165E%3D&response-content-disposition=attachment%3Bfilename%3D%22hello2.txt%22&x-amz-security-token=some_amz_security_token"
+        },
+        {
+            "file": "keton/elo.txt",
+            "url": "https://ps-projects.s3.amazonaws.com/some/path/artifacts/keton/elo.txt?AWSAccessKeyId=some_aws_access_key_id&Expires=713274132&Signature=tHriojGx03S%2FKkVGQGVI5CQRFTo%3D&response-content-disposition=attachment%3Bfilename%3D%22elo.txt%22&x-amz-security-token=some_amz_security_token"
+        }
+    ],
+}
 
 GET_PRESIGNED_URL_FOR_S3_BUCKET_RESPONSE_JSON = {
     "data": {

tests/functional/test_jobs.py

@@ -394,14 +394,14 @@ class TestListJobArtifacts(TestJobs):
 
     @mock.patch("gradient.api_sdk.clients.http_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()
+        get_patched.return_value = MockResponse(LIST_JOB_FILES_RESPONSE_JSON)
         job_id = "some_job_id"
         result = self.runner.invoke(cli.cli,
                                     ["jobs", "artifacts", "list", "--id", job_id, "--apiKey", "some_key", "--size",
                                      "--links",
                                      "--files", "foo"])
 
-        get_patched.assert_called_with("{}/jobs/artifactsList".format(self.URL),
+        get_patched.assert_called_with("{}/jobs/artifactsListV2".format(self.URL),
                                        headers=EXPECTED_HEADERS_WITH_CHANGED_API_KEY,
                                        json=None,
                                        params={"jobId": job_id,
@@ -412,11 +412,11 @@ def test_should_send_valid_get_request_with_all_parameters_for_a_list_of_artifac
 
     @mock.patch("gradient.api_sdk.clients.http_client.requests.get")
     def test_should_read_options_from_yaml_file(self, get_patched, jobs_artifacts_list_config_path):
-        get_patched.return_value = MockResponse()
+        get_patched.return_value = MockResponse(LIST_JOB_FILES_RESPONSE_JSON)
         command = ["jobs", "artifacts", "list", "--optionsFile", jobs_artifacts_list_config_path]
         result = self.runner.invoke(cli.cli, command)
 
-        get_patched.assert_called_with("{}/jobs/artifactsList".format(self.URL),
+        get_patched.assert_called_with("{}/jobs/artifactsListV2".format(self.URL),
                                        headers=EXPECTED_HEADERS_WITH_CHANGED_API_KEY,
                                        json=None,
                                        params={"files": "keton*.py",
@@ -434,12 +434,12 @@ def test_should_send_valid_get_request_with_valid_param_for_a_list_of_artifacts_
                                                                                                               get_patched,
                                                                                                               option,
                                                                                                               param):
-        get_patched.return_value = MockResponse(status_code=200)
+        get_patched.return_value = MockResponse(LIST_JOB_FILES_RESPONSE_JSON, status_code=200)
         job_id = "some_job_id"
         result = self.runner.invoke(cli.cli,
                                     ["jobs", "artifacts", "list", "--id", job_id, "--apiKey", "some_key"] + [option])
 
-        get_patched.assert_called_with("{}/jobs/artifactsList".format(self.URL),
+        get_patched.assert_called_with("{}/jobs/artifactsListV2".format(self.URL),
                                        headers=EXPECTED_HEADERS_WITH_CHANGED_API_KEY,
                                        json=None,
                                        params={"jobId": job_id,
@@ -635,7 +635,7 @@ def test_should_send_proper_data_and_tag_job(self, post_patched, get_patched, pu
 
 class TestDownloadJobArtifacts(TestJobs):
     runner = CliRunner()
-    LIST_FILES_URL = "https://api.paperspace.io/jobs/artifactsList"
+    LIST_FILES_URL = "https://api.paperspace.io/jobs/artifactsListV2"
     DESTINATION_DIR_NAME = "dest"
     DESTINATION_DIR_PATH = os.path.join(tempfile.gettempdir(), "dest")