From 037d07cea0bf0cafbb77f10183a33c4f9cc22134 Mon Sep 17 00:00:00 2001 From: Jonathan Feng Date: Wed, 17 Jan 2024 03:02:48 -0800 Subject: [PATCH 1/5] add task batch metadata and task templates support --- scaleapi/__init__.py | 39 ++++++++++++++++++++++++++++++++++++++- scaleapi/_version.py | 2 +- scaleapi/batches.py | 1 + scaleapi/projects.py | 31 +++++++++++++++++++++++++++++++ tests/test_client.py | 16 ++++++++++++---- 5 files changed, 83 insertions(+), 6 deletions(-) diff --git a/scaleapi/__init__.py b/scaleapi/__init__.py index ae96c64..12185bb 100644 --- a/scaleapi/__init__.py +++ b/scaleapi/__init__.py @@ -4,7 +4,7 @@ from scaleapi.evaluation_tasks import EvaluationTask from scaleapi.exceptions import ScaleInvalidRequest from scaleapi.files import File -from scaleapi.projects import Project +from scaleapi.projects import Project, TaskTemplate from scaleapi.training_tasks import TrainingTask from ._version import __version__ # noqa: F401 @@ -619,6 +619,7 @@ def create_batch( callback: str = "", calibration_batch: bool = False, self_label_batch: bool = False, + metadata: Dict = {}, ) -> Batch: """Create a new Batch within a project. https://docs.scale.com/reference#batch-creation @@ -639,6 +640,8 @@ def create_batch( Only applicable for self serve projects. Create a self_label batch by setting the self_label_batch flag to true. + metadata (Dict): + Optional metadata to be stored at the TaskBatch level Returns: Batch: Created batch object @@ -650,6 +653,7 @@ def create_batch( calibration_batch=calibration_batch, self_label_batch=self_label_batch, callback=callback, + metadata=metadata, ) batchdata = self.api.post_request(endpoint, body=payload) return Batch(batchdata, self) @@ -828,6 +832,22 @@ def get_batches( offset += batches.limit has_more = batches.has_more + def set_batch_metadata(self, batch_name: str, metadata: Dict) -> Batch: + """Sets metadata for a TaskBatch. + + Args: + batch_name (str): + Batch name + metadata (Dict): + Metadata to set for TaskBatch + + Returns: + Batch + """ + endpoint = f"batches/{Api.quote_string(batch_name)}/setMetadata" + batchdata = self.api.post_request(endpoint, body=metadata) + return Batch(batchdata, self) + def create_project( self, project_name: str, @@ -933,6 +953,23 @@ def update_project(self, project_name: str, **kwargs) -> Project: endpoint = f"projects/{Api.quote_string(project_name)}/setParams" projectdata = self.api.post_request(endpoint, body=kwargs) return Project(projectdata, self) + + def get_project_template(self, project_name: str) -> Dict: + """Gets the task template of a project if a template exists. + Throws an error if the project task-type does not support Task Templates. + Currently only TextCollection and Chat task types support Task Templates. + + Args: + project_name (str): + Project's name + + Returns: + Template | None + """ + endpoint = f"projects/{Api.quote_string(project_name)}/taskTemplates" + template = self.api.get_request(endpoint) + return TaskTemplate(template, self) + def upload_file(self, file: IO, **kwargs) -> File: """Upload file. diff --git a/scaleapi/_version.py b/scaleapi/_version.py index d6ce145..f60d06a 100644 --- a/scaleapi/_version.py +++ b/scaleapi/_version.py @@ -1,2 +1,2 @@ -__version__ = "2.15.2" +__version__ = "2.15.3" __package_name__ = "scaleapi" diff --git a/scaleapi/batches.py b/scaleapi/batches.py index d0cb73a..6a0ffc3 100644 --- a/scaleapi/batches.py +++ b/scaleapi/batches.py @@ -19,6 +19,7 @@ def __init__(self, json, client): self.project = json["project"] self.created_at = json["created_at"] self.project = json["project"] + self.metadata = json['metadata'] self.tasks_pending = None self.tasks_completed = None diff --git a/scaleapi/projects.py b/scaleapi/projects.py index 8a0536f..e5d2d9c 100644 --- a/scaleapi/projects.py +++ b/scaleapi/projects.py @@ -1,3 +1,30 @@ +class TaskTemplate: + """Task Template Object.""" + def __init__(self, json, client): + self._json = json + self._client = client + self.id = json['id'] + self.project = json['project'] + self.version = json['version'] + self.created_at = json['created_at'] + self.updated_at = json['updated_at'] + self.template_variables = json['template_variables'] + + def __hash__(self): + return hash(self.id) + + def __str__(self): + return f"TaskTemplate(id={self.id}, project={self.project})" + + def __repr__(self): + return f"TaskTemplate({self._json})" + + def get_template_variables(self): + return self.template_variables + + def as_dict(self): + return self._json + class Project: """Project class, containing Project information.""" @@ -24,6 +51,10 @@ def __str__(self): def __repr__(self): return f"Project({self._json})" + + def get_template(self) -> TaskTemplate: + """Returns TaskTemplate. Only works for Chat and TextCollection type.""" + return self._client.get_project_template(self.name) def as_dict(self): """Returns all attributes as a dictionary""" diff --git a/tests/test_client.py b/tests/test_client.py index 63d4c67..fe67181 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -1,5 +1,5 @@ # pylint: disable=missing-function-docstring - +import os import time import uuid from datetime import datetime @@ -21,7 +21,7 @@ try: print(f"SDK Version: {scaleapi.__version__}") - test_api_key = "test_fe79860cdbe547bf91b4e7da897a6c92" + test_api_key = os.environ["SCALE_TEST_API_KEY"] if test_api_key.startswith("test_") or test_api_key.endswith("|test"): client = scaleapi.ScaleClient(test_api_key, "pytest") @@ -392,6 +392,7 @@ def create_a_batch(): callback="http://www.example.com/callback", batch_name=str(uuid.uuid4()), project=TEST_PROJECT_NAME, + metadata={'some_key': 'some_value'}, ) @@ -432,12 +433,14 @@ def test_get_batch_status(): assert batch2.status == BatchStatus.InProgress.value + def test_get_batch(): batch = create_a_batch() batch2 = client.get_batch(batch.name) assert batch.name == batch2.name assert batch2.status == BatchStatus.InProgress.value - + # test metadata + assert batch2.metadata['some_key'] == 'some_value' def test_batches(): batches = [] @@ -458,6 +461,11 @@ def test_get_batches(): all_batches = list(client.get_batches(project_name=TEST_PROJECT_NAME)) assert total_batches == len(all_batches) +def test_set_batch_metadata(): + batch = create_a_batch() + batch = client.set_batch_metadata(batch.name, {'new_key': 'new_value'}) + assert batch.metadata['new_key'] == 'new_value' + def test_files_upload(): with open("tests/test_image.png", "rb") as f: @@ -496,7 +504,7 @@ def test_invite_teammates(): try: project = client.get_project(STUDIO_TEST_PROJECT) except ScaleResourceNotFound: - client.create_project(project_name=STUDIO_TEST_PROJECT) + client.create_project(project_name=STUDIO_TEST_PROJECT, task_type=TaskType.ImageAnnotation) STUDIO_BATCH_TEST_NAME = f"studio-test-batch-{current_timestamp}" From 186b1e43f4a7f8ea2d88baa7c7628b07fa52843c Mon Sep 17 00:00:00 2001 From: Jonathan Feng Date: Wed, 17 Jan 2024 13:14:54 -0800 Subject: [PATCH 2/5] reformatting --- scaleapi/__init__.py | 11 +++++------ scaleapi/batches.py | 2 +- scaleapi/projects.py | 22 ++++++++++++---------- tests/test_client.py | 15 +++++++++------ 4 files changed, 27 insertions(+), 23 deletions(-) diff --git a/scaleapi/__init__.py b/scaleapi/__init__.py index 12185bb..db1e195 100644 --- a/scaleapi/__init__.py +++ b/scaleapi/__init__.py @@ -834,13 +834,13 @@ def get_batches( def set_batch_metadata(self, batch_name: str, metadata: Dict) -> Batch: """Sets metadata for a TaskBatch. - + Args: batch_name (str): Batch name metadata (Dict): Metadata to set for TaskBatch - + Returns: Batch """ @@ -953,16 +953,16 @@ def update_project(self, project_name: str, **kwargs) -> Project: endpoint = f"projects/{Api.quote_string(project_name)}/setParams" projectdata = self.api.post_request(endpoint, body=kwargs) return Project(projectdata, self) - + def get_project_template(self, project_name: str) -> Dict: """Gets the task template of a project if a template exists. Throws an error if the project task-type does not support Task Templates. Currently only TextCollection and Chat task types support Task Templates. - + Args: project_name (str): Project's name - + Returns: Template | None """ @@ -970,7 +970,6 @@ def get_project_template(self, project_name: str) -> Dict: template = self.api.get_request(endpoint) return TaskTemplate(template, self) - def upload_file(self, file: IO, **kwargs) -> File: """Upload file. Refer to Files API Reference: diff --git a/scaleapi/batches.py b/scaleapi/batches.py index 6a0ffc3..743b2bc 100644 --- a/scaleapi/batches.py +++ b/scaleapi/batches.py @@ -19,7 +19,7 @@ def __init__(self, json, client): self.project = json["project"] self.created_at = json["created_at"] self.project = json["project"] - self.metadata = json['metadata'] + self.metadata = json["metadata"] self.tasks_pending = None self.tasks_completed = None diff --git a/scaleapi/projects.py b/scaleapi/projects.py index e5d2d9c..af576a3 100644 --- a/scaleapi/projects.py +++ b/scaleapi/projects.py @@ -1,15 +1,16 @@ class TaskTemplate: """Task Template Object.""" + def __init__(self, json, client): self._json = json self._client = client - self.id = json['id'] - self.project = json['project'] - self.version = json['version'] - self.created_at = json['created_at'] - self.updated_at = json['updated_at'] - self.template_variables = json['template_variables'] - + self.id = json["id"] + self.project = json["project"] + self.version = json["version"] + self.created_at = json["created_at"] + self.updated_at = json["updated_at"] + self.template_variables = json["template_variables"] + def __hash__(self): return hash(self.id) @@ -18,13 +19,14 @@ def __str__(self): def __repr__(self): return f"TaskTemplate({self._json})" - + def get_template_variables(self): return self.template_variables - + def as_dict(self): return self._json + class Project: """Project class, containing Project information.""" @@ -51,7 +53,7 @@ def __str__(self): def __repr__(self): return f"Project({self._json})" - + def get_template(self) -> TaskTemplate: """Returns TaskTemplate. Only works for Chat and TextCollection type.""" return self._client.get_project_template(self.name) diff --git a/tests/test_client.py b/tests/test_client.py index fe67181..9bf52ea 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -392,7 +392,7 @@ def create_a_batch(): callback="http://www.example.com/callback", batch_name=str(uuid.uuid4()), project=TEST_PROJECT_NAME, - metadata={'some_key': 'some_value'}, + metadata={"some_key": "some_value"}, ) @@ -433,14 +433,14 @@ def test_get_batch_status(): assert batch2.status == BatchStatus.InProgress.value - def test_get_batch(): batch = create_a_batch() batch2 = client.get_batch(batch.name) assert batch.name == batch2.name assert batch2.status == BatchStatus.InProgress.value # test metadata - assert batch2.metadata['some_key'] == 'some_value' + assert batch2.metadata["some_key"] == "some_value" + def test_batches(): batches = [] @@ -461,10 +461,11 @@ def test_get_batches(): all_batches = list(client.get_batches(project_name=TEST_PROJECT_NAME)) assert total_batches == len(all_batches) + def test_set_batch_metadata(): batch = create_a_batch() - batch = client.set_batch_metadata(batch.name, {'new_key': 'new_value'}) - assert batch.metadata['new_key'] == 'new_value' + batch = client.set_batch_metadata(batch.name, {"new_key": "new_value"}) + assert batch.metadata["new_key"] == "new_value" def test_files_upload(): @@ -504,7 +505,9 @@ def test_invite_teammates(): try: project = client.get_project(STUDIO_TEST_PROJECT) except ScaleResourceNotFound: - client.create_project(project_name=STUDIO_TEST_PROJECT, task_type=TaskType.ImageAnnotation) + client.create_project( + project_name=STUDIO_TEST_PROJECT, task_type=TaskType.ImageAnnotation + ) STUDIO_BATCH_TEST_NAME = f"studio-test-batch-{current_timestamp}" From 58ea4eb64a236439f4d0a0d1067a5da3a03a5a74 Mon Sep 17 00:00:00 2001 From: Jonathan Feng Date: Wed, 17 Jan 2024 13:16:32 -0800 Subject: [PATCH 3/5] type fixes --- scaleapi/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/scaleapi/__init__.py b/scaleapi/__init__.py index db1e195..e589370 100644 --- a/scaleapi/__init__.py +++ b/scaleapi/__init__.py @@ -954,7 +954,7 @@ def update_project(self, project_name: str, **kwargs) -> Project: projectdata = self.api.post_request(endpoint, body=kwargs) return Project(projectdata, self) - def get_project_template(self, project_name: str) -> Dict: + def get_project_template(self, project_name: str) -> TaskTemplate: """Gets the task template of a project if a template exists. Throws an error if the project task-type does not support Task Templates. Currently only TextCollection and Chat task types support Task Templates. @@ -964,7 +964,7 @@ def get_project_template(self, project_name: str) -> Dict: Project's name Returns: - Template | None + TaskTemplate """ endpoint = f"projects/{Api.quote_string(project_name)}/taskTemplates" template = self.api.get_request(endpoint) From 693d3f14629cdd9286532632c10eefddbe83bb33 Mon Sep 17 00:00:00 2001 From: Jonathan Feng Date: Wed, 17 Jan 2024 13:18:28 -0800 Subject: [PATCH 4/5] comment length fix --- scaleapi/__init__.py | 5 +++-- scaleapi/projects.py | 3 ++- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/scaleapi/__init__.py b/scaleapi/__init__.py index e589370..83123b8 100644 --- a/scaleapi/__init__.py +++ b/scaleapi/__init__.py @@ -956,8 +956,9 @@ def update_project(self, project_name: str, **kwargs) -> Project: def get_project_template(self, project_name: str) -> TaskTemplate: """Gets the task template of a project if a template exists. - Throws an error if the project task-type does not support Task Templates. - Currently only TextCollection and Chat task types support Task Templates. + Throws an error if the project task-type does not support + Task Templates. Currently only TextCollection and Chat task + types support Task Templates. Args: project_name (str): diff --git a/scaleapi/projects.py b/scaleapi/projects.py index af576a3..f7094e1 100644 --- a/scaleapi/projects.py +++ b/scaleapi/projects.py @@ -55,7 +55,8 @@ def __repr__(self): return f"Project({self._json})" def get_template(self) -> TaskTemplate: - """Returns TaskTemplate. Only works for Chat and TextCollection type.""" + """Returns TaskTemplate. + Only works for Chat and TextCollection type.""" return self._client.get_project_template(self.name) def as_dict(self): From f36c1569b55430d8cf9c8b6614ee978b3fd3083c Mon Sep 17 00:00:00 2001 From: Jonathan Feng Date: Wed, 17 Jan 2024 14:17:04 -0800 Subject: [PATCH 5/5] docstrings --- scaleapi/__init__.py | 4 ++-- scaleapi/projects.py | 2 ++ 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/scaleapi/__init__.py b/scaleapi/__init__.py index 83123b8..26523a7 100644 --- a/scaleapi/__init__.py +++ b/scaleapi/__init__.py @@ -619,7 +619,7 @@ def create_batch( callback: str = "", calibration_batch: bool = False, self_label_batch: bool = False, - metadata: Dict = {}, + metadata: Dict = None, ) -> Batch: """Create a new Batch within a project. https://docs.scale.com/reference#batch-creation @@ -653,7 +653,7 @@ def create_batch( calibration_batch=calibration_batch, self_label_batch=self_label_batch, callback=callback, - metadata=metadata, + metadata=metadata or {}, ) batchdata = self.api.post_request(endpoint, body=payload) return Batch(batchdata, self) diff --git a/scaleapi/projects.py b/scaleapi/projects.py index f7094e1..17147d6 100644 --- a/scaleapi/projects.py +++ b/scaleapi/projects.py @@ -21,9 +21,11 @@ def __repr__(self): return f"TaskTemplate({self._json})" def get_template_variables(self): + """Returns template variables dictionary""" return self.template_variables def as_dict(self): + """Returns task template object as JSON dictionary""" return self._json