improved docker base

2024-07-17 12:32:39 +08:00 · 2024-07-17 12:32:39 +08:00 · 59c57f983a
parent 08fb1c5671
commit 59c57f983a
2 changed files with 4 additions and 0 deletions
--- a/build_scripts/inference/serve
+++ b/build_scripts/inference/serve
--- a/middleware_api/trainings/create_training_job.py
+++ b/middleware_api/trainings/create_training_job.py
@ -137,6 +137,10 @@ def _trigger_sagemaker_training_job(
            "s3_location": f"s3://{bucket_name}/{train_params_file}",
        },
        job_id=train_job.id,
+        environment={
+            "SITE_PACKAGES_S3_PATH": f"aws-gcr-solutions-{region}/"
+                                     f"stable-diffusion-aws-extension-github-mainline/{esd_version}/train.tar"
+        }
    )
    est.fit(wait=False)