create archive when env yaml is provided in publish command

ads/opctl/conda/cmds.py

@@ -46,6 +46,7 @@
 from ads.opctl.config.base import ConfigProcessor
 from ads.opctl.config.merger import ConfigMerger
 from ads.opctl.conda.multipart_uploader import MultiPartUploader
+import tempfile
 
 
 def _fetch_manifest_template() -> Dict:
@@ -108,6 +109,7 @@ def _create(
  conda_pack_folder: str,
  gpu: bool,
  overwrite: bool,
+ prepare_publish: bool = False,
 ) -> str:
  """Create a conda pack given an environment yaml file under conda pack folder specified.
 
@@ -123,6 +125,8 @@ def _create(
  whether to build against GPU image
  overwrite : bool
  whether to overwrite existing pack of the same slug
+ prepare_pubish : bool
+ whether to create conda pack archive after conda pack is created
 
  Raises
  ------
@@ -180,6 +184,11 @@ def _create(
  manifest["manifest"]["manifest_version"] = "1.0"
 
  logger.info(f"Creating conda environment {slug}")
+ conda_dep = None
+ with open(env_file) as mfile:
+ conda_dep = yaml.safe_load(mfile.read())
+ conda_dep["manifest"] = manifest["manifest"]
+
  if is_in_notebook_session() or NO_CONTAINER:
  command = f"conda env create --prefix {pack_folder_path} --file {os.path.abspath(os.path.expanduser(env_file))}"
  run_command(command, shell=True)
@@ -191,35 +200,56 @@ def _create(
  )
 
  create_command = f"conda env create --prefix {docker_pack_folder_path} --file {docker_env_file_path}"
-
+ 
  volumes = {
  pack_folder_path: {"bind": docker_pack_folder_path},
  os.path.abspath(os.path.expanduser(env_file)): {
  "bind": docker_env_file_path
  },
+
  }
+
  if gpu:
  image = ML_JOB_GPU_IMAGE
  else:
  image = ML_JOB_IMAGE
  try:
- run_container(
- image=image, bind_volumes=volumes, env_vars={}, command=create_command
- )
+ if prepare_publish:
+ tmp_file = tempfile.NamedTemporaryFile(suffix=".yaml")
+ # Save the manifest in the temp file that can be mounted inside the container so that archiving will work
+ with open(tmp_file.name, 'w') as f:
+ yaml.safe_dump(conda_dep, f) 
+
+ pack_script = os.path.join(os.path.dirname(os.path.abspath(__file__)), "pack.py")
+ 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')}"
+
+ # add pack script and manifest file to the mount so that archive can be created in the same container run
+ condapack_script = {
+ pack_script: {"bind": os.path.join(DEFAULT_IMAGE_HOME_DIR, "pack.py")},
+ tmp_file.name: {"bind": os.path.join(DEFAULT_IMAGE_HOME_DIR, "manifest.yaml")}
+ }
+ volumes = {**volumes, **condapack_script} # | not supported in python 3.8
+
+ run_container(
+ image=image, bind_volumes=volumes, entrypoint="/bin/bash -c ", env_vars={}, command=f" '{create_command} && {pack_command}'"
+ )
+ else:
+ run_container(
+ image=image, bind_volumes=volumes, env_vars={}, command=create_command
+ ) 
  except Exception:
  if os.path.exists(pack_folder_path):
  shutil.rmtree(pack_folder_path)
  raise RuntimeError(f"Could not create environment {slug}.")
 
- conda_dep = None
- with open(env_file) as mfile:
- conda_dep = yaml.safe_load(mfile.read())
- conda_dep["manifest"] = manifest["manifest"]
- with open(f"{os.path.join(pack_folder_path, slug)}_manifest.yaml", "w") as mfile:
+ # Save the manifest file inside the host machine, where the conda environment is saved.
+ manifest_location = f"{os.path.join(pack_folder_path, slug)}_manifest.yaml"
+ with open(manifest_location, "w") as mfile:
  yaml.safe_dump(conda_dep, mfile)
 
  logger.info(f"Environment `{slug}` setup complete.")
  print(f"Pack {slug} created under {pack_folder_path}.")
+
  return slug
 
 
@@ -467,6 +497,7 @@ def _install(
 def publish(**kwargs) -> None:
  p = ConfigProcessor().step(ConfigMerger, **kwargs)
  exec_config = p.config["execution"]
+ skip_archive = False
  if exec_config.get("environment_file", None):
  name = _get_name(exec_config.get("name"), exec_config.get("environment_file"))
  slug = _create(
@@ -476,7 +507,9 @@ def publish(**kwargs) -> None:
  conda_pack_folder=exec_config["conda_pack_folder"],
  gpu=exec_config.get("gpu", False),
  overwrite=exec_config["overwrite"],
+ prepare_publish=True
  )
+ skip_archive = True # The conda pack archive is already created during create process.
  else:
  slug = exec_config.get("slug")
  if not slug:
@@ -493,9 +526,10 @@ def publish(**kwargs) -> None:
  oci_profile=exec_config.get("oci_profile"),
  overwrite=exec_config["overwrite"],
  auth_type=exec_config["auth"],
+ skip_archive=skip_archive
  )
 
-
+ 
 def _publish(
  conda_slug: str,
  conda_uri_prefix: str,
@@ -504,6 +538,7 @@ def _publish(
  oci_profile: str,
  overwrite: bool,
  auth_type: str,
+ skip_archive: bool = False
 ) -> None:
  """Publish a local conda pack to object storage location
 
@@ -579,29 +614,30 @@ def _publish(
  publish_slug = "_".join(ans.lower().split(" "))
 
  pack_script = os.path.join(os.path.dirname(os.path.abspath(__file__)), "pack.py")
- if is_in_notebook_session() or NO_CONTAINER:
- command = f"python {pack_script} {pack_folder_path}"
- run_command(command, shell=True)
- else:
- volumes = {
- pack_folder_path: {
- "bind": os.path.join(DEFAULT_IMAGE_HOME_DIR, conda_slug)
- },
- pack_script: {"bind": os.path.join(DEFAULT_IMAGE_HOME_DIR, "pack.py")},
- }
- command = f"python {os.path.join(DEFAULT_IMAGE_HOME_DIR, 'pack.py')} {os.path.join(DEFAULT_IMAGE_HOME_DIR, conda_slug)}"
- gpu = env["manifest"]["arch_type"] == "GPU"
- _check_job_image_exists(gpu)
- if gpu:
- image = ML_JOB_GPU_IMAGE
+ if not skip_archive:
+ if is_in_notebook_session() or NO_CONTAINER:
+ command = f"python {pack_script} {pack_folder_path}"
+ run_command(command, shell=True)
  else:
- image = ML_JOB_IMAGE
- try:
- run_container(
- image=image, bind_volumes=volumes, env_vars={}, command=command
- )
- except Exception:
- raise RuntimeError(f"Could not pack environment {conda_slug}.")
+ volumes = {
+ pack_folder_path: {
+ "bind": os.path.join(DEFAULT_IMAGE_HOME_DIR, conda_slug)
+ },
+ pack_script: {"bind": os.path.join(DEFAULT_IMAGE_HOME_DIR, "pack.py")},
+ }
+ command = f"python {os.path.join(DEFAULT_IMAGE_HOME_DIR, 'pack.py')} {os.path.join(DEFAULT_IMAGE_HOME_DIR, conda_slug)}"
+ gpu = env["manifest"]["arch_type"] == "GPU"
+ _check_job_image_exists(gpu)
+ if gpu:
+ image = ML_JOB_GPU_IMAGE
+ else:
+ image = ML_JOB_IMAGE
+ try:
+ run_container(
+ image=image, bind_volumes=volumes, env_vars={}, command=command
+ )
+ except Exception:
+ raise RuntimeError(f"Could not pack environment {conda_slug}.")
 
  pack_file = os.path.join(pack_folder_path, f"{conda_slug}.tar.gz")
  if not os.path.exists(pack_file):

ads/opctl/conda/pack.py

@@ -16,11 +16,14 @@
 import conda_pack
 
 import yaml
+import argparse
 
 
-def main(pack_folder_path):
+def main(pack_folder_path, manifest_file=None):
  slug = os.path.basename(pack_folder_path)
- manifest_path = glob.glob(os.path.join(pack_folder_path, "*_manifest.yaml"))[0]
+ manifest_path = (
+ manifest_file or glob.glob(os.path.join(pack_folder_path, "*_manifest.yaml"))[0]
+ )
  with open(manifest_path) as f:
  env = yaml.safe_load(f.read())
 
@@ -59,8 +62,10 @@ def main(pack_folder_path):
  raise RuntimeError(
  "Error creating the pack file using `conda_pack.pack()`."
  )
+ print(f"Copy {pack_file} to {pack_folder_path}")
  shutil.copy(pack_file, pack_folder_path)
  file_path = os.path.join(pack_folder_path, os.path.basename(pack_file))
+ print(f"Pack built at {file_path}")
  print(
  f"changing permission for {file_path}",
  flush=True,
@@ -69,4 +74,14 @@ def main(pack_folder_path):
 
 
 if __name__ == "__main__":
- main(sys.argv[1])
+ parser = argparse.ArgumentParser(
+ prog="Prepare conda archive",
+ description="Uses conda_pack library to pack the conda environment.",
+ )
+ parser.add_argument("--conda-path", type=str, help="Path to the conda environment")
+ parser.add_argument(
+ "--manifest-location", type=str, default=None, help="Path to manifest location"
+ )
+ args = parser.parse_args()
+
+ main(args.conda_path, args.manifest_location)