add support for mlflow in caliban shell

caliban/docker/build.py

@@ -125,6 +125,9 @@
   recommend using an absolute path!
 
   """
+  if from_path is None:
+    return ""
+
   cmd = f"COPY --chown={user_id}:{user_group} {from_path} {to_path}\n"
 
   if comment is not None:
@@ -358,6 +361,7 @@
   caliban_config = caliban_config or {}
 
   arg = package.main_module or package.script_path
+  arg = [arg] if arg is not None else []
   package_path = package.package_path
 
   copy_code = copy_command(
@@ -369,7 +373,7 @@
 
   # This needs to use json so that quotes print as double quotes, not single
   # quotes.
-  executable_s = json.dumps(package.executable + [arg])
+  executable_s = json.dumps(package.executable + arg)
 
   entrypoint_code = _generate_entrypoint(
       executable=executable_s,

caliban/platform/run.py

@@ -328,9 +328,23 @@
   if image_id is None:
     image_id = b.build_image(job_mode, **build_image_kwargs)
 
+  caliban_config = build_image_kwargs.get('caliban_config', {})
+
   base_cmd = _run_cmd(job_mode, run_args)
 
-  command = base_cmd + [image_id] + script_args
+  launcher_args = um.mlflow_args(
+      caliban_config=caliban_config,
+      experiment_name=um.mlflow_shell_experiment_name(),
+      index=-1,
+      tags={
+          um.GPU_ENABLED_TAG: str(job_mode == c.JobMode.GPU).lower(),
+          um.TPU_ENABLED_TAG: 'false',
+          um.DOCKER_IMAGE_TAG: image_id,
+          um.PLATFORM_TAG: Platform.LOCAL.value,
+      },
+  )
+
+  command = base_cmd + [image_id] + launcher_args + script_args
 
   logging.info("Running command: {}".format(' '.join(command)))
   subprocess.call(command)

caliban/platform/shell.py

@@ -25,6 +25,7 @@
 import caliban.config as c
 import caliban.docker.build as b
 import caliban.platform.run as r
+import caliban.util as u
 
 
 def _home_mount_cmds(enable_home_mount: bool) -> List[str]:
@@ -99,9 +100,13 @@
   if entrypoint is None:
     entrypoint = b.SHELL_DICT[shell].executable
 
+  build_image_kwargs['package'] = u.Package(executable=[entrypoint],
+                                            script_path=None,
+                                            main_module=None,
+                                            package_path=None)
+
   interactive_run_args = _interactive_opts(workdir) + [
-      "-it", \
-      "--entrypoint", entrypoint
+      "-it"
   ] + _home_mount_cmds(mount_home) + run_args
 
   r.run(job_mode=job_mode,

caliban/util/metrics.py

@@ -176,7 +176,7 @@ def mlflow_args(
 
   caliban_config: caliban configuration dict
   experiment: experiment object
-  index: job index
+  index: job index, if < 0, then no run name is generated
   tags: dictionary of tags to pass to mlflow
 
   Returns:
@@ -188,6 +188,14 @@ def mlflow_args(
 
   env = {f'ENVVAR_{k}': v for k, v in tags.items()}
   env['MLFLOW_EXPERIMENT_NAME'] = experiment_name
-  env['MLFLOW_RUN_NAME'] = _mlflow_job_name(index=index)
+
+  if index >= 0:
+    env['MLFLOW_RUN_NAME'] = _mlflow_job_name(index=index)
 
   return ['--caliban_config', json.dumps({'env': env})]
+
+
+def mlflow_shell_experiment_name() -> str:
+  '''generates an experiment name for a caliban shell session'''
+  timestamp = datetime.datetime.now().strftime('%Y-%m-%d-%H-%M-%S')
+  return f'{u.current_user()}-shell-{timestamp}'

tests/caliban/util/test_metrics.py

@@ -221,3 +221,9 @@ def test_launcher_config_file():
 
     # make sure we clean up appropriately
     assert not os.path.exists(cfg_fname)
+
+
+def test_mlflow_shell_experiment_name():
+  experiment_name = um.mlflow_shell_experiment_name()
+  assert experiment_name is not None
+  assert len(experiment_name) > 0