Launch v1 compatibility (#387)

Adds Launch v1 client compatibility to the Nucleus Launch integration

Author: seanshi-scale

CHANGELOG.md

@@ -5,6 +5,11 @@ All notable changes to the [Nucleus Python Client](https://github.com/scaleapi/n
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.15.9](https://github.com/scaleapi/nucleus-python-client/releases/tag/v0.15.9) - 2023-06-26
+
+### Added
+- Support for Scale Launch client v1.0.0 and higher for the Nucleus + Launch integration
+
 ## [0.15.7](https://github.com/scaleapi/nucleus-python-client/releases/tag/v0.15.7) - 2023-06-09
 
 ### Added

nucleus/__init__.py

@@ -592,6 +592,9 @@ def create_launch_model(
 
         A model bundle consists of exactly {predict_fn_or_cls}, {load_predict_fn + model}, or {load_predict_fn + load_model_fn}.
         Pre/post-processing code can be included inside load_predict_fn/model or in predict_fn_or_cls call.
+        Note: the exact parameters used will depend on the version of the Launch client used.
+        i.e. if you are on Launch client version 0.x, you will use `env_params`, otherwise
+        you will use `pytorch_image_tag` and `tensorflow_version`.
 
         Parameters:
             model_bundle_name: Name of model bundle you want to create. This acts as a unique identifier.
@@ -608,7 +611,8 @@ def create_launch_model(
                 ["tensorflow==2.3.0", "tensorflow-hub==0.11.0"]. If no list has been passed, will default to the currently
                 imported list of packages.
             app_config: Either a Dictionary that represents a YAML file contents or a local path to a YAML file.
-            env_params: A dictionary that dictates environment information e.g.
+            env_params: Only for launch v0.
+                A dictionary that dictates environment information e.g.
                 the use of pytorch or tensorflow, which cuda/cudnn versions to use.
                 Specifically, the dictionary should contain the following keys:
                 "framework_type": either "tensorflow" or "pytorch".
@@ -617,6 +621,9 @@ def create_launch_model(
                 "cudnn_version" Version of cudnn used, e.g. "cudnn8-devel".
                 "tensorflow_version": Version of tensorflow, e.g. "2.3.0". Only applicable if framework_type is tensorflow
             globals_copy: Dictionary of the global symbol table. Normally provided by `globals()` built-in function.
+            pytorch_image_tag: Only for launch v1, and if you want to use pytorch framework type.
+                The tag of the pytorch docker image you want to use, e.g. 1.11.0-cuda11.3-cudnn8-runtime
+            tensorflow_version: Only for launch v1, and if you want to use tensorflow. Version of tensorflow, e.g. "2.3.0".
         """
         from launch import LaunchClient
 
@@ -637,13 +644,22 @@ def create_launch_model(
             "model_bundle_name": name + "-nucleus-autogen",
             **bundle_args,
         }
-
-        bundle = launch_client.create_model_bundle(**kwargs)
+        if hasattr(launch_client, "create_model_bundle_from_callable_v2"):
+            # Launch client is >= 1.0.0
+            bundle = launch_client.create_model_bundle_from_callable_v2(
+                **kwargs
+            )
+            bundle_name = (
+                bundle.name
+            )  # both v0 and v1 have a .name field but are different types
+        else:
+            bundle = launch_client.create_model_bundle(**kwargs)
+            bundle_name = bundle.name
         return self.create_model(
             name,
             reference_id,
             metadata,
-            bundle.name,
+            bundle_name,
         )
 
     def create_launch_model_from_dir(
@@ -705,12 +721,16 @@ def create_launch_model_from_dir(
         as the desired inference loading function, then the `load_predict_fn_module_path` argument should be
         `my_module1.my_inference_file.f`.
 
+        Note: the exact keys for `bundle_from_dir_args` used will depend on the version of the Launch client used.
+        i.e. if you are on Launch client version 0.x, you will use `env_params`, otherwise
+        you will use `pytorch_image_tag` and `tensorflow_version`.
 
         Keys for `bundle_from_dir_args`:
             model_bundle_name: Name of model bundle you want to create. This acts as a unique identifier.
             base_paths: The paths on the local filesystem where the bundle code lives.
             requirements_path: A path on the local filesystem where a requirements.txt file lives.
-            env_params: A dictionary that dictates environment information e.g.
+            env_params: Only for launch v0.
+                A dictionary that dictates environment information e.g.
                 the use of pytorch or tensorflow, which cuda/cudnn versions to use.
                 Specifically, the dictionary should contain the following keys:
                 "framework_type": either "tensorflow" or "pytorch".
@@ -723,6 +743,9 @@ def create_launch_model_from_dir(
             load_model_fn_module_path: A python module path for a function that returns a model. The output feeds into
                 the function located at load_predict_fn_module_path.
             app_config: Either a Dictionary that represents a YAML file contents or a local path to a YAML file.
+            pytorch_image_tag: Only for launch v1, and if you want to use pytorch framework type.
+                The tag of the pytorch docker image you want to use, e.g. 1.11.0-cuda11.3-cudnn8-runtime
+            tensorflow_version: Only for launch v1, and if you want to use tensorflow. Version of tensorflow, e.g. "2.3.0".
         """
         from launch import LaunchClient
 
@@ -744,13 +767,21 @@ def create_launch_model_from_dir(
             **bundle_from_dir_args,
         }
 
-        bundle = launch_client.create_model_bundle_from_dirs(**kwargs)
+        if hasattr(launch_client, "create_model_bundle_from_dirs_v2"):
+            # Launch client is >= 1.0.0, use new fn
+            bundle = launch_client.create_model_bundle_from_dirs_v2(**kwargs)
+            # Different code paths give different types for bundle, although both have a .name field
+            bundle_name = bundle.name
+        else:
+            # Launch client is < 1.0.0
+            bundle = launch_client.create_model_bundle_from_dirs(**kwargs)
+            bundle_name = bundle.name
 
         return self.create_model(
             name,
             reference_id,
             metadata,
-            bundle.name,
+            bundle_name,
         )
 
     @deprecated(

nucleus/dataset_item_uploader.py

@@ -163,7 +163,6 @@ def get_form_data_and_file_pointers_fn(
         """
 
         def fn():
-
             # For some reason, our backend only accepts this reformatting of items when
             # doing local upload.
             # TODO: make it just accept the same exact format as a normal append request

nucleus/metrics/segmentation_metrics.py

@@ -195,7 +195,8 @@ def _calculate_confusion_matrix(
 
     def _is_instance_segmentation(self, annotation, prediction):
         """Guesses that we're dealing with instance segmentation if we have multiple segments with same label.
-        Degenerate case is same as semseg so fine to misclassify in that case."""
+        Degenerate case is same as semseg so fine to misclassify in that case.
+        """
         # This is a trick to get ordered sets
         ann_labels = list(
             dict.fromkeys(s.label for s in annotation.annotations)
@@ -525,7 +526,6 @@ def _metric_impl(
         annotation: SegmentationAnnotation,
         prediction: SegmentationPrediction,
     ) -> ScalarResult:
-
         ap_per_threshold = []
         weight = 0
         for iou_threshold in self.iou_thresholds:

pyproject.toml

@@ -21,7 +21,7 @@ exclude = '''
 
 [tool.poetry]
 name = "scale-nucleus"
-version = "0.15.8"
+version = "0.15.9"
 description = "The official Python client library for Nucleus, the Data Platform for AI"
 license =  "MIT"
 authors = ["Scale AI Nucleus Team <nucleusapi@scaleapi.com>"]

tests/test_models.py

@@ -144,7 +144,6 @@ def testing_model(ref_id):
 
 
 def test_remove_invalid_tag_from_model(CLIENT, dataset: Dataset):
-
     model_reference = "model_" + str(time.time())
     model = CLIENT.create_model(TEST_MODEL_NAME, model_reference)
 

tests/test_scene.py

@@ -268,7 +268,6 @@ def test_scene_add_frame():
 
 def test_video_scene_property_methods():
     for scene_json in TEST_VIDEO_SCENES["scenes"]:
-
         scene = VideoScene.from_json(scene_json)
 
         expected_frame_rate = scene_json.get("frame_rate", None)

tests/test_segmentation.py

@@ -83,7 +83,6 @@ def test_single_local_semseg_gt_upload(dataset: Dataset):
 
 
 def test_batch_local_semseg_gt_upload(dataset: Dataset):
-
     # This reference id is not in the dataset.
     bad_reference_id = TEST_SEGMENTATION_ANNOTATIONS[-1]["reference_id"]
 

tests/test_utils.py

@@ -11,7 +11,6 @@ def weird_function():
 
 
 def test_serialize():
-
     test_items = [
         DatasetItem("fake_url1", "fake_id1"),
         DatasetItem(