chore(tests): refactor root_cache with default True; pass bucket name

Author: gadicc

tests/integration/lib.py

@@ -169,7 +169,7 @@ def getMinio(id="disposable"):
             print("Reusing existing global minio")
             aws_access_key_id = os.getenv("AWS_ACCESS_KEY_ID")
             aws_secret_access_key = os.getenv("AWS_SECRET_ACCESS_KEY")
-
+            aws_s3_default_bucket = AWS_S3_DEFAULT_BUCKET
             s3 = boto3.client(
                 "s3",
                 endpoint_url=endpoint_url,
@@ -187,6 +187,7 @@ def getMinio(id="disposable"):
                 "endpoint_url": endpoint_url,
                 "aws_access_key_id": aws_access_key_id,
                 "aws_secret_access_key": aws_secret_access_key,
+                "aws_s3_default_bucket": aws_s3_default_bucket,
                 "s3": s3,
             }
             _minioCache.update({id: result})
@@ -222,6 +223,7 @@ def onstop():
 
     aws_access_key_id = "minioadmin"
     aws_secret_access_key = "minioadmin"
+    aws_s3_default_bucket = AWS_S3_DEFAULT_BUCKET
     s3 = boto3.client(
         "s3",
         endpoint_url=endpoint_url,
@@ -241,6 +243,7 @@ def onstop():
         "endpoint_url": endpoint_url,
         "aws_access_key_id": aws_access_key_id,
         "aws_secret_access_key": aws_secret_access_key,
+        "aws_s3_default_bucket": aws_s3_default_bucket,
         "s3": s3,
     }
     _minioCache.update({id: result})
@@ -251,7 +254,13 @@ def onstop():
 
 
 def getDDA(
-    minio=None, command=None, environment={}, stream_logs=False, wait=True, **kwargs
+    minio=None,
+    command=None,
+    environment={},
+    stream_logs=False,
+    wait=True,
+    root_cache=True,
+    **kwargs,
 ):
     global _ddaCache
     if _ddaCache:
@@ -278,7 +287,7 @@ def getDDA(
                 "AWS_ACCESS_KEY_ID": minio.aws_access_key_id,
                 "AWS_SECRET_ACCESS_KEY": minio.aws_secret_access_key,
                 "AWS_DEFAULT_REGION": "",
-                "AWS_S3_DEFAULT_BUCKET": AWS_S3_DEFAULT_BUCKET,
+                "AWS_S3_DEFAULT_BUCKET": minio.aws_s3_default_bucket,
                 "AWS_S3_ENDPOINT_URL": minio.endpoint_url,
             }
         )
@@ -287,13 +296,16 @@ def onstop():
         global _ddaCache
         _ddaCache = None
 
+    HOME = os.getenv("HOME")
+
     container, stop = startContainer(
         "gadicc/diffusers-api:test",
         command,
         stream_logs=stream_logs,
         ports={8000: port},
         device_requests=[docker.types.DeviceRequest(count=-1, capabilities=[["gpu"]])],
         environment=environment,
+        volumes=root_cache and [f"{HOME}/root-cache:/root/.cache"],
         onstop=onstop,
         **kwargs,
     )

tests/integration/test_build_download.py

@@ -2,33 +2,39 @@
 from .lib import getMinio, getDDA
 from test import runTest
 
+
 def test_cloudcache_build_download():
     """
     Download a model from cloud-cache at build time (no HuggingFace)
     """
     minio = getMinio()
     print(minio)
     environment = {
-      "RUNTIME_DOWNLOADS": 0,
-      "MODEL_ID": "stabilityai/stable-diffusion-2-1-base",
-      "MODEL_PRECISION": "fp16",
-      "MODEL_REVISION": "fp16",
-      "MODEL_URL": "s3://", # <--
+        "RUNTIME_DOWNLOADS": 0,
+        "MODEL_ID": "stabilityai/stable-diffusion-2-1-base",
+        "MODEL_PRECISION": "fp16",
+        "MODEL_REVISION": "fp16",
+        "MODEL_URL": "s3://",  # <--
     }
-    conda="conda run --no-capture-output -n xformers"
+    conda = "conda run --no-capture-output -n xformers"
     dda = getDDA(
-      minio=minio,
-      stream_logs=True,
-      environment=environment,
-      command=["sh", "-c", f"{conda} python3 -u download.py && ls -l && {conda} python3 -u server.py"],
+        minio=minio,
+        stream_logs=True,
+        environment=environment,
+        root_cache=False,
+        command=[
+            "sh",
+            "-c",
+            f"{conda} python3 -u download.py && ls -l && {conda} python3 -u server.py",
+        ],
     )
     print(dda)
     assert dda.container.status == "running"
 
     ## bucket.objects.all().delete()
     result = runTest(
         "txt2img",
-        { "test_url": dda.url },
+        {"test_url": dda.url},
         {
             "MODEL_ID": "stabilityai/stable-diffusion-2-1-base",
         },
@@ -48,30 +54,35 @@ def test_huggingface_build_download():
     uploads if missing.
     """
     environment = {
-      "RUNTIME_DOWNLOADS": 0,
-      "MODEL_ID": "stabilityai/stable-diffusion-2-1-base",
-      "MODEL_PRECISION": "fp16",
-      "MODEL_REVISION": "fp16",
+        "RUNTIME_DOWNLOADS": 0,
+        "MODEL_ID": "stabilityai/stable-diffusion-2-1-base",
+        "MODEL_PRECISION": "fp16",
+        "MODEL_REVISION": "fp16",
     }
-    conda="conda run --no-capture-output -n xformers"
+    conda = "conda run --no-capture-output -n xformers"
     dda = getDDA(
-      stream_logs=True,
-      environment=environment,
-      command=["sh", "-c", f"{conda} python3 -u download.py && ls -l && {conda} python3 -u server.py"],
+        stream_logs=True,
+        environment=environment,
+        root_cache=False,
+        command=[
+            "sh",
+            "-c",
+            f"{conda} python3 -u download.py && ls -l && {conda} python3 -u server.py",
+        ],
     )
     print(dda)
     assert dda.container.status == "running"
 
     ## bucket.objects.all().delete()
     result = runTest(
         "txt2img",
-        { "test_url": dda.url },
+        {"test_url": dda.url},
         {
             "MODEL_ID": "stabilityai/stable-diffusion-2-1-base",
             # "MODEL_ID": "hf-internal-testing/tiny-stable-diffusion-pipe",
             "MODEL_PRECISION": "fp16",
             "MODEL_REVISION": "fp16",
-            "MODEL_URL": "", # <-- no model_url, i.e. no cloud cache
+            "MODEL_URL": "",  # <-- no model_url, i.e. no cloud cache
         },
         {"num_inference_steps": 1},
     )

tests/integration/test_cloud_cache.py

@@ -2,20 +2,21 @@
 from .lib import getMinio, getDDA
 from test import runTest
 
+
 def test_cloud_cache_create_and_upload():
     """
     Check if model exists in cloud cache bucket download otherwise, save
     with safetensors, and upload model.tar.zst to bucket
     """
     minio = getMinio()
     print(minio)
-    dda = getDDA(minio=minio, stream_logs=True)
+    dda = getDDA(minio=minio, stream_logs=True, root_cache=False)
     print(dda)
 
     ## bucket.objects.all().delete()
     result = runTest(
         "txt2img",
-        { "test_url": dda.url },
+        {"test_url": dda.url},
         {
             "MODEL_ID": "stabilityai/stable-diffusion-2-1-base",
             # "MODEL_ID": "hf-internal-testing/tiny-stable-diffusion-pipe",

tests/integration/test_general.py

@@ -23,11 +23,9 @@ def setup_class(self):
         print("setup_class")
         # self.minio = minio = getMinio()
 
-        HOME = os.getenv("HOME")
         self.dda = dda = getDDA(
             # minio=minio
             # stream_logs=True,
-            volumes=[f"{HOME}/root-cache:/root/.cache"],
         )
         print(dda)
 

tests/integration/test_memory.py

@@ -9,11 +9,9 @@ def test_memory():
     Make sure when switching models we release VRAM afterwards.
     """
     minio = getMinio("global")
-    HOME = os.getenv("HOME")
     dda = getDDA(
         minio=minio,
         stream_logs=True,
-        volumes=[f"{HOME}/root-cache:/root/.cache"],
     )
     print(dda)