use custom dlc

autogluon · Jun 8, 2024 · 7d9fda5 · 7d9fda5
1 parent 55c3169
commit 7d9fda5
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 4 deletions.
diff --git a/src/autogluon/cloud/backend/ray_backend.py b/src/autogluon/cloud/backend/ray_backend.py
@@ -395,6 +395,10 @@ def _get_image_uri(self, framework_version: str, instance_type: str, custom_imag
                 image_scope="training",
                 instance_type=instance_type,
             )
+            if "g4dn" in instance_type:
+                image_uri = "369469875935.dkr.ecr.us-east-1.amazonaws.com/weisu:gpu-latest"
+            else:
+                image_uri = "369469875935.dkr.ecr.us-east-1.amazonaws.com/weisu:cpu-latest"
         return image_uri
 
     def _construct_ag_args(self, predictor_init_args, predictor_fit_args, **kwargs):

diff --git a/src/autogluon/cloud/utils/ag_sagemaker.py b/src/autogluon/cloud/utils/ag_sagemaker.py
@@ -39,6 +39,10 @@ def __init__(
                 image_scope="training",
                 instance_type=instance_type,
             )
+            if "g4dn" in instance_type:
+                self.image_uri = "369469875935.dkr.ecr.us-east-1.amazonaws.com/weisu:gpu-latest"
+            else:
+                self.image_uri = "369469875935.dkr.ecr.us-east-1.amazonaws.com/weisu:cpu-latest"
         super().__init__(
             entry_point=entry_point,
             source_dir=source_dir,
@@ -74,6 +78,10 @@ def create_model(
             image_scope="inference",
             instance_type=instance_type,
         )
+        if instance_type == "gpu":
+            image_uri = "369469875935.dkr.ecr.us-east-1.amazonaws.com/weisu:gpu-inference-latest"
+        else:
+            image_uri = "369469875935.dkr.ecr.us-east-1.amazonaws.com/weisu:cpu-inference-latest"
         if predictor_cls is None:
 
             def predict_wrapper(endpoint, session):
@@ -141,6 +149,10 @@ def __init__(
                 image_scope="inference",
                 instance_type=instance_type,
             )
+        if instance_type == "gpu":
+            image_uri = "369469875935.dkr.ecr.us-east-1.amazonaws.com/weisu:gpu-inference-latest"
+        else:
+            image_uri = "369469875935.dkr.ecr.us-east-1.amazonaws.com/weisu:cpu-inference-latest"
         # setting PYTHONUNBUFFERED to disable output buffering for endpoints logging
         if env is None:
             env = {}

diff --git a/tests/conftest.py b/tests/conftest.py
@@ -10,10 +10,10 @@
 
 
 class CloudTestHelper:
-    cpu_training_image = "369469875935.dkr.ecr.us-east-1.amazonaws.com/autogluon-nightly-training:cpu-latest"
-    gpu_training_image = "369469875935.dkr.ecr.us-east-1.amazonaws.com/autogluon-nightly-training:gpu-latest"
-    cpu_inference_image = "369469875935.dkr.ecr.us-east-1.amazonaws.com/autogluon-nightly-inference:cpu-latest"
-    gpu_inference_image = "369469875935.dkr.ecr.us-east-1.amazonaws.com/autogluon-nightly-inference:gpu-latest"
+    cpu_training_image = "369469875935.dkr.ecr.us-east-1.amazonaws.com/weisu:cpu-latest"
+    gpu_training_image = "369469875935.dkr.ecr.us-east-1.amazonaws.com/weisu:gpu-latest"
+    cpu_inference_image = "369469875935.dkr.ecr.us-east-1.amazonaws.com/weisu:cpu-inference-latest"
+    gpu_inference_image = "369469875935.dkr.ecr.us-east-1.amazonaws.com/weisu:gpu-inference-latest"
 
     @staticmethod
     def get_custom_image_uri(framework_version="source", type="training", gpu=False):