46
46
from ads .opctl .config .base import ConfigProcessor
47
47
from ads .opctl .config .merger import ConfigMerger
48
48
from ads .opctl .conda .multipart_uploader import MultiPartUploader
49
+ import tempfile
49
50
50
51
51
52
def _fetch_manifest_template () -> Dict :
@@ -108,6 +109,7 @@ def _create(
108
109
conda_pack_folder : str ,
109
110
gpu : bool ,
110
111
overwrite : bool ,
112
+ prepare_publish : bool = False ,
111
113
) -> str :
112
114
"""Create a conda pack given an environment yaml file under conda pack folder specified.
113
115
@@ -123,6 +125,8 @@ def _create(
123
125
whether to build against GPU image
124
126
overwrite : bool
125
127
whether to overwrite existing pack of the same slug
128
+ prepare_pubish : bool
129
+ whether to create conda pack archive after conda pack is created
126
130
127
131
Raises
128
132
------
@@ -180,6 +184,11 @@ def _create(
180
184
manifest ["manifest" ]["manifest_version" ] = "1.0"
181
185
182
186
logger .info (f"Creating conda environment { slug } " )
187
+ conda_dep = None
188
+ with open (env_file ) as mfile :
189
+ conda_dep = yaml .safe_load (mfile .read ())
190
+ conda_dep ["manifest" ] = manifest ["manifest" ]
191
+
183
192
if is_in_notebook_session () or NO_CONTAINER :
184
193
command = f"conda env create --prefix { pack_folder_path } --file { os .path .abspath (os .path .expanduser (env_file ))} "
185
194
run_command (command , shell = True )
@@ -191,35 +200,56 @@ def _create(
191
200
)
192
201
193
202
create_command = f"conda env create --prefix { docker_pack_folder_path } --file { docker_env_file_path } "
194
-
203
+
195
204
volumes = {
196
205
pack_folder_path : {"bind" : docker_pack_folder_path },
197
206
os .path .abspath (os .path .expanduser (env_file )): {
198
207
"bind" : docker_env_file_path
199
208
},
209
+
200
210
}
211
+
201
212
if gpu :
202
213
image = ML_JOB_GPU_IMAGE
203
214
else :
204
215
image = ML_JOB_IMAGE
205
216
try :
206
- run_container (
207
- image = image , bind_volumes = volumes , env_vars = {}, command = create_command
208
- )
217
+ if prepare_publish :
218
+ tmp_file = tempfile .NamedTemporaryFile (suffix = ".yaml" )
219
+ # Save the manifest in the temp file that can be mounted inside the container so that archiving will work
220
+ with open (tmp_file .name , 'w' ) as f :
221
+ yaml .safe_dump (conda_dep , f )
222
+
223
+ pack_script = os .path .join (os .path .dirname (os .path .abspath (__file__ )), "pack.py" )
224
+ pack_command = f"python { os .path .join (DEFAULT_IMAGE_HOME_DIR , 'pack.py' )} --conda-path { docker_pack_folder_path } --manifest-location { os .path .join (DEFAULT_IMAGE_HOME_DIR , 'manifest.yaml' )} "
225
+
226
+ # add pack script and manifest file to the mount so that archive can be created in the same container run
227
+ condapack_script = {
228
+ pack_script : {"bind" : os .path .join (DEFAULT_IMAGE_HOME_DIR , "pack.py" )},
229
+ tmp_file .name : {"bind" : os .path .join (DEFAULT_IMAGE_HOME_DIR , "manifest.yaml" )}
230
+ }
231
+ volumes = {** volumes , ** condapack_script } # | not supported in python 3.8
232
+
233
+ run_container (
234
+ image = image , bind_volumes = volumes , entrypoint = "/bin/bash -c " , env_vars = {}, command = f" '{ create_command } && { pack_command } '"
235
+ )
236
+ else :
237
+ run_container (
238
+ image = image , bind_volumes = volumes , env_vars = {}, command = create_command
239
+ )
209
240
except Exception :
210
241
if os .path .exists (pack_folder_path ):
211
242
shutil .rmtree (pack_folder_path )
212
243
raise RuntimeError (f"Could not create environment { slug } ." )
213
244
214
- conda_dep = None
215
- with open (env_file ) as mfile :
216
- conda_dep = yaml .safe_load (mfile .read ())
217
- conda_dep ["manifest" ] = manifest ["manifest" ]
218
- with open (f"{ os .path .join (pack_folder_path , slug )} _manifest.yaml" , "w" ) as mfile :
245
+ # Save the manifest file inside the host machine, where the conda environment is saved.
246
+ manifest_location = f"{ os .path .join (pack_folder_path , slug )} _manifest.yaml"
247
+ with open (manifest_location , "w" ) as mfile :
219
248
yaml .safe_dump (conda_dep , mfile )
220
249
221
250
logger .info (f"Environment `{ slug } ` setup complete." )
222
251
print (f"Pack { slug } created under { pack_folder_path } ." )
252
+
223
253
return slug
224
254
225
255
@@ -467,6 +497,7 @@ def _install(
467
497
def publish (** kwargs ) -> None :
468
498
p = ConfigProcessor ().step (ConfigMerger , ** kwargs )
469
499
exec_config = p .config ["execution" ]
500
+ skip_archive = False
470
501
if exec_config .get ("environment_file" , None ):
471
502
name = _get_name (exec_config .get ("name" ), exec_config .get ("environment_file" ))
472
503
slug = _create (
@@ -476,7 +507,9 @@ def publish(**kwargs) -> None:
476
507
conda_pack_folder = exec_config ["conda_pack_folder" ],
477
508
gpu = exec_config .get ("gpu" , False ),
478
509
overwrite = exec_config ["overwrite" ],
510
+ prepare_publish = True
479
511
)
512
+ skip_archive = True # The conda pack archive is already created during create process.
480
513
else :
481
514
slug = exec_config .get ("slug" )
482
515
if not slug :
@@ -493,9 +526,10 @@ def publish(**kwargs) -> None:
493
526
oci_profile = exec_config .get ("oci_profile" ),
494
527
overwrite = exec_config ["overwrite" ],
495
528
auth_type = exec_config ["auth" ],
529
+ skip_archive = skip_archive
496
530
)
497
531
498
-
532
+
499
533
def _publish (
500
534
conda_slug : str ,
501
535
conda_uri_prefix : str ,
@@ -504,6 +538,7 @@ def _publish(
504
538
oci_profile : str ,
505
539
overwrite : bool ,
506
540
auth_type : str ,
541
+ skip_archive : bool = False
507
542
) -> None :
508
543
"""Publish a local conda pack to object storage location
509
544
@@ -579,29 +614,30 @@ def _publish(
579
614
publish_slug = "_" .join (ans .lower ().split (" " ))
580
615
581
616
pack_script = os .path .join (os .path .dirname (os .path .abspath (__file__ )), "pack.py" )
582
- if is_in_notebook_session () or NO_CONTAINER :
583
- command = f"python { pack_script } { pack_folder_path } "
584
- run_command (command , shell = True )
585
- else :
586
- volumes = {
587
- pack_folder_path : {
588
- "bind" : os .path .join (DEFAULT_IMAGE_HOME_DIR , conda_slug )
589
- },
590
- pack_script : {"bind" : os .path .join (DEFAULT_IMAGE_HOME_DIR , "pack.py" )},
591
- }
592
- command = f"python { os .path .join (DEFAULT_IMAGE_HOME_DIR , 'pack.py' )} { os .path .join (DEFAULT_IMAGE_HOME_DIR , conda_slug )} "
593
- gpu = env ["manifest" ]["arch_type" ] == "GPU"
594
- _check_job_image_exists (gpu )
595
- if gpu :
596
- image = ML_JOB_GPU_IMAGE
617
+ if not skip_archive :
618
+ if is_in_notebook_session () or NO_CONTAINER :
619
+ command = f"python { pack_script } --conda-path { pack_folder_path } "
620
+ run_command (command , shell = True )
597
621
else :
598
- image = ML_JOB_IMAGE
599
- try :
600
- run_container (
601
- image = image , bind_volumes = volumes , env_vars = {}, command = command
602
- )
603
- except Exception :
604
- raise RuntimeError (f"Could not pack environment { conda_slug } ." )
622
+ volumes = {
623
+ pack_folder_path : {
624
+ "bind" : os .path .join (DEFAULT_IMAGE_HOME_DIR , conda_slug )
625
+ },
626
+ pack_script : {"bind" : os .path .join (DEFAULT_IMAGE_HOME_DIR , "pack.py" )},
627
+ }
628
+ command = f"python { os .path .join (DEFAULT_IMAGE_HOME_DIR , 'pack.py' )} --conda-path { os .path .join (DEFAULT_IMAGE_HOME_DIR , conda_slug )} "
629
+ gpu = env ["manifest" ]["arch_type" ] == "GPU"
630
+ _check_job_image_exists (gpu )
631
+ if gpu :
632
+ image = ML_JOB_GPU_IMAGE
633
+ else :
634
+ image = ML_JOB_IMAGE
635
+ try :
636
+ run_container (
637
+ image = image , bind_volumes = volumes , env_vars = {}, command = command
638
+ )
639
+ except Exception :
640
+ raise RuntimeError (f"Could not pack environment { conda_slug } ." )
605
641
606
642
pack_file = os .path .join (pack_folder_path , f"{ conda_slug } .tar.gz" )
607
643
if not os .path .exists (pack_file ):
0 commit comments