chore: import upstream snapshot with attribution

mlflow/java/client/src/main/java/org/mlflow/api/proto/package-info.java

@@ -0,0 +1,2 @@
+/** Autogenerated MLflow Tracking API entity objects. */
+package org.mlflow.api.proto;

mlflow/java/client/src/main/java/org/mlflow/artifacts/ArtifactRepository.java

@@ -0,0 +1,113 @@
+package org.mlflow.artifacts;
+
+import java.io.File;
+import java.util.List;
+
+import org.mlflow.api.proto.Service.FileInfo;
+
+/**
+ * Allows logging, listing, and downloading artifacts against a remote Artifact Repository.
+ * This is used for storing potentially-large objects associated with MLflow runs.
+ */
+public interface ArtifactRepository {
+
+  /**
+   * Uploads the given local file to the run's root artifact directory. For example,
+   *
+   *   <pre>
+   *   logArtifact("/my/localModel")
+   *   listArtifacts() // returns "localModel"
+   *   </pre>
+   *
+   * @param localFile File to upload. Must exist, and must be a simple file (not a directory).
+   */
+  void logArtifact(File localFile);
+
+  /**
+   * Uploads the given local file to an artifactPath within the run's root directory. For example,
+   *
+   *   <pre>
+   *   logArtifact("/my/localModel", "model")
+   *   listArtifacts("model") // returns "model/localModel"
+   *   </pre>
+   *
+   * (i.e., the localModel file is now available in model/localModel).
+   *
+   * @param localFile File to upload. Must exist, and must be a simple file (not a directory).
+   * @param artifactPath Artifact path relative to the run's root directory. Should NOT
+   *                     start with a /.
+   */
+  void logArtifact(File localFile, String artifactPath);
+
+  /**
+   * Uploads all files within the given local director the run's root artifact directory.
+   * For example, if /my/local/dir/ contains two files "file1" and "file2", then
+   *
+   *   <pre>
+   *   logArtifacts("/my/local/dir")
+   *   listArtifacts() // returns "file1" and "file2"
+   *   </pre>
+   *
+   * @param localDir Directory to upload. Must exist, and must be a directory (not a simple file).
+   */
+  void logArtifacts(File localDir);
+
+
+  /**
+   * Uploads all files within the given local director an artifactPath within the run's root
+   * artifact directory. For example, if /my/local/dir/ contains two files "file1" and "file2", then
+   *
+   *   <pre>
+   *   logArtifacts("/my/local/dir", "model")
+   *   listArtifacts("model") // returns "model/file1" and "model/file2"
+   *   </pre>
+   *
+   * (i.e., the contents of the local directory are now available in model/).
+   *
+   * @param localDir Directory to upload. Must exist, and must be a directory (not a simple file).
+   * @param artifactPath Artifact path relative to the run's root directory. Should NOT
+   *                     start with a /.
+   */
+  void logArtifacts(File localDir, String artifactPath);
+
+  /**
+   * Lists the artifacts immediately under the run's root artifact directory. This does not
+   * recursively list; instead, it will return FileInfos with isDir=true where further
+   * listing may be done.
+   */
+  List<FileInfo> listArtifacts();
+
+  /**
+   * Lists the artifacts immediately under the given artifactPath within the run's root artifact
+   * irectory. This does not recursively list; instead, it will return FileInfos with isDir=true
+   * where further listing may be done.
+   * @param artifactPath Artifact path relative to the run's root directory. Should NOT
+   *                     start with a /.
+   */
+  List<FileInfo> listArtifacts(String artifactPath);
+
+  /**
+   * Returns a local directory containing *all* artifacts within the run's artifact directory.
+   * Note that this will download the entire directory path, and so may be expensive if
+   * the directory a lot of data.
+   */
+  File downloadArtifacts();
+
+  /**
+   * Returns a local file or directory containing all artifacts within the given artifactPath
+   * within the run's root artifactDirectory. For example, if "model/file1" and "model/file2"
+   * exist within the artifact directory, then
+   *
+   *   <pre>
+   *   downloadArtifacts("model") // returns a local directory containing "file1" and "file2"
+   *   downloadArtifacts("model/file1") // returns a local *file* with the contents of file1.
+   *   </pre>
+   *
+   * Note that this will download the entire subdirectory path, and so may be expensive if
+   * the subdirectory a lot of data.
+   *
+   * @param artifactPath Artifact path relative to the run's root directory. Should NOT
+   *                     start with a /.
+   */
+  File downloadArtifacts(String artifactPath);
+}

mlflow/java/client/src/main/java/org/mlflow/artifacts/ArtifactRepositoryFactory.java

@@ -0,0 +1,17 @@
+package org.mlflow.artifacts;
+
+import java.net.URI;
+
+import org.mlflow.tracking.creds.MlflowHostCredsProvider;
+
+public class ArtifactRepositoryFactory {
+  private final MlflowHostCredsProvider hostCredsProvider;
+
+  public ArtifactRepositoryFactory(MlflowHostCredsProvider hostCredsProvider) {
+    this.hostCredsProvider = hostCredsProvider;
+  }
+
+  public ArtifactRepository getArtifactRepository(URI baseArtifactUri, String runId) {
+    return new CliBasedArtifactRepository(baseArtifactUri.toString(), runId, hostCredsProvider);
+  }
+}

mlflow/java/client/src/main/java/org/mlflow/artifacts/CliBasedArtifactRepository.java

@@ -0,0 +1,314 @@
+package org.mlflow.artifacts;
+
+import java.io.File;
+import java.io.IOException;
+import java.lang.reflect.Type;
+import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.Lists;
+import com.google.gson.Gson;
+import com.google.gson.reflect.TypeToken;
+import com.google.protobuf.InvalidProtocolBufferException;
+import com.google.protobuf.util.JsonFormat;
+import org.apache.commons.io.IOUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.mlflow.api.proto.Service;
+import org.mlflow.tracking.MlflowClientException;
+import org.mlflow.tracking.creds.MlflowHostCreds;
+import org.mlflow.tracking.creds.DatabricksMlflowHostCreds;
+import org.mlflow.tracking.creds.MlflowHostCredsProvider;
+
+/**
+ * Shells out to the 'mlflow' command line utility to upload, download, and list artifacts. This
+ * is used as a fallback to implement any artifact repositories which are not natively supported
+ * within Java.
+ *
+ * We require that 'mlflow' is available in the system path.
+ */
+public class CliBasedArtifactRepository implements ArtifactRepository {
+  private static final Logger logger = LoggerFactory.getLogger(CliBasedArtifactRepository.class);
+
+  // Global check if we ever successfully loaded 'mlflow'. This allows us to print a more
+  // helpful error message if the executable is not in the path.
+  private static final AtomicBoolean mlflowSuccessfullyLoaded = new AtomicBoolean(false);
+
+  // Name of the Python CLI utility which can be exec'd directly, with MLflow on its path
+  private final String PYTHON_EXECUTABLE =
+    Optional.ofNullable(System.getenv("MLFLOW_PYTHON_EXECUTABLE")).orElse("python");
+
+  // Python CLI command
+  private final String PYTHON_COMMAND = "mlflow.store.artifact.cli";
+
+  // Base directory of the artifactory, used to let the user know why this repository was chosen.
+  private final String artifactBaseDir;
+
+  // Run ID this repository is targeting.
+  private final String runId;
+
+  // Used to pass credentials as environment variables
+  // (e.g., MLFLOW_TRACKING_URI or DATABRICKS_HOST) to the mlflow process.
+  private final MlflowHostCredsProvider hostCredsProvider;
+
+  public CliBasedArtifactRepository(
+      String artifactBaseDir,
+      String runId,
+      MlflowHostCredsProvider hostCredsProvider) {
+    this.artifactBaseDir = artifactBaseDir;
+    this.runId = runId;
+    this.hostCredsProvider = hostCredsProvider;
+  }
+
+  @Override
+  public void logArtifact(File localFile, String artifactPath) {
+    checkMlflowAccessible();
+    if (!localFile.exists()) {
+      throw new MlflowClientException("Local file does not exist: " + localFile);
+    }
+    if (localFile.isDirectory()) {
+      throw new MlflowClientException("Local path points to a directory. Use logArtifacts" +
+        " instead: " + localFile);
+    }
+
+    List<String> baseCommand = Lists.newArrayList(
+      "log-artifact", "--local-file", localFile.toString());
+    List<String> command = appendRunIdArtifactPath(baseCommand, runId, artifactPath);
+    String tag = "log file " + localFile + " to " + getTargetIdentifier(artifactPath);
+    forkMlflowProcess(command, tag);
+  }
+
+  @Override
+  public void logArtifact(File localFile) {
+    logArtifact(localFile, null);
+  }
+
+  @Override
+  public void logArtifacts(File localDir, String artifactPath) {
+    checkMlflowAccessible();
+    if (!localDir.exists()) {
+      throw new MlflowClientException("Local file does not exist: " + localDir);
+    }
+    if (localDir.isFile()) {
+      throw new MlflowClientException("Local path points to a file. Use logArtifact" +
+        " instead: " + localDir);
+    }
+
+    List<String> baseCommand = Lists.newArrayList(
+      "log-artifacts", "--local-dir", localDir.toString());
+    List<String> command = appendRunIdArtifactPath(baseCommand, runId, artifactPath);
+    String tag = "log dir " + localDir + " to " + getTargetIdentifier(artifactPath);
+    forkMlflowProcess(command, tag);
+  }
+
+  @Override
+  public void logArtifacts(File localDir) {
+    logArtifacts(localDir, null);
+  }
+
+  @Override
+  public File downloadArtifacts(String artifactPath) {
+    checkMlflowAccessible();
+    String tag = "download artifacts for " + getTargetIdentifier(artifactPath);
+    List<String> command = appendRunIdArtifactPath(
+      Lists.newArrayList("download"), runId, artifactPath);
+    String stdOutput = forkMlflowProcess(command, tag);
+    String[] splits = stdOutput.split(System.lineSeparator());
+    return new File(splits[splits.length-1].trim());
+  }
+
+  @Override
+  public File downloadArtifacts() {
+    return downloadArtifacts(null);
+  }
+
+  @Override
+  public List<Service.FileInfo> listArtifacts(String artifactPath) {
+    checkMlflowAccessible();
+    String tag = "list artifacts in " + getTargetIdentifier(artifactPath);
+    List<String> command = appendRunIdArtifactPath(
+      Lists.newArrayList("list"), runId, artifactPath);
+    String jsonOutput = forkMlflowProcess(command, tag);
+    return parseFileInfos(jsonOutput);
+  }
+
+  @Override
+  public List<Service.FileInfo> listArtifacts() {
+    return listArtifacts(null);
+  }
+
+  /**
+   * Only available in the CliBasedArtifactRepository. Downloads an artifact to the local
+   * filesystem when provided with an artifact uri. This method should not be used directly
+   * by the user. Please use {@link org.mlflow.tracking.MlflowClient}
+   *
+   * @param artifactUri Artifact uri
+   * @return Directory/file of the artifact
+   */
+  public File downloadArtifactFromUri(String artifactUri) {
+    checkMlflowAccessible();
+    String tag = "download artifacts for " + artifactUri;
+    List<String> command = Lists.newArrayList("download", "--artifact-uri", artifactUri);
+    String localPath = forkMlflowProcess(command, tag).trim();
+    return new File(localPath);
+  }
+
+  /** Parses a list of JSON FileInfos, as returned by 'mlflow artifacts list'. */
+  private List<Service.FileInfo> parseFileInfos(String json) {
+    // The protobuf deserializer doesn't allow us to directly deserialize a list, so we
+    // deserialize a list-of-dictionaries, and then reserialize each dictionary to pass it to
+    // the protobuf deserializer.
+    Gson gson = new Gson();
+    Type type = new TypeToken<List<Map<String, Object>>>(){}.getType();
+    List<Map<String, Object>> listOfDicts = gson.fromJson(json, type);
+    List<Service.FileInfo> fileInfos = new ArrayList<>();
+    for (Map<String, Object> dict: listOfDicts) {
+      String fileInfoJson = gson.toJson(dict);
+      try {
+        Service.FileInfo.Builder builder = Service.FileInfo.newBuilder();
+        JsonFormat.parser().merge(fileInfoJson, builder);
+        fileInfos.add(builder.build());
+      } catch (InvalidProtocolBufferException e) {
+        throw new MlflowClientException("Failed to deserialize JSON into FileInfo: " + json, e);
+      }
+    }
+    return fileInfos;
+  }
+
+  /**
+   * Checks whether the 'mlflow' executable is available, and throws a nice error if not.
+   * If this method has ever run successfully before (in the entire JVM), we will not rerun it.
+   */
+  private void checkMlflowAccessible() {
+    if (mlflowSuccessfullyLoaded.get()) {
+      return;
+    }
+
+    try {
+      String tag = "get mlflow version";
+      forkMlflowProcess(Lists.newArrayList("--help"), tag);
+      logger.info("Found local mlflow executable");
+      mlflowSuccessfullyLoaded.set(true);
+    } catch (MlflowClientException e) {
+      String errorMessage = String.format("Failed to exec '%s -m %s', needed to" +
+          " access artifacts within the non-Java-native artifact store at '%s'. Please make" +
+          " sure mlflow is available on your local system path (e.g., from 'pip install mlflow')",
+        PYTHON_EXECUTABLE, PYTHON_COMMAND, artifactBaseDir);
+      throw new MlflowClientException(errorMessage, e);
+    }
+  }
+
+  /**
+   * Forks the given mlflow command and awaits for its successful completion.
+   *
+   * @param mlflowCommand List of arguments to invoke mlflow with.
+   * @param tag User-facing tag which will be used to identify what we were trying to do
+   *            in the case of a failure.
+   * @return raw stdout of the process, decoded as a utf-8 string
+   * @throws MlflowClientException if the process exits with a non-zero exit code, or anything
+   *                               else goes wrong.
+   */
+  private String forkMlflowProcess(List<String> mlflowCommand, String tag) {
+    String stdout;
+    Process process = null;
+    try {
+      MlflowHostCreds hostCreds = hostCredsProvider.getHostCreds();
+      List<String> fullCommand = Lists.newArrayList(PYTHON_EXECUTABLE, "-m", PYTHON_COMMAND);
+      fullCommand.addAll(mlflowCommand);
+      ProcessBuilder pb = new ProcessBuilder(fullCommand);
+      if (hostCreds instanceof DatabricksMlflowHostCreds) {
+        setProcessEnvironmentDatabricks(pb.environment(), (DatabricksMlflowHostCreds) hostCreds);
+      } else {
+        setProcessEnvironment(pb.environment(), hostCreds);
+      }
+      process = pb.start();
+      stdout = IOUtils.toString(process.getInputStream(), StandardCharsets.UTF_8);
+      int exitValue = process.waitFor();
+      if (exitValue != 0) {
+        throw new MlflowClientException("Failed to " + tag + ". Error: " +
+          getErrorBestEffort(process));
+      }
+    } catch (IOException | InterruptedException e) {
+      throw new MlflowClientException("Failed to fork mlflow process to " + tag +
+        ". Process stderr: " + getErrorBestEffort(process), e);
+    }
+    return stdout;
+  }
+
+  @VisibleForTesting
+  void setProcessEnvironment(Map<String, String> environment, MlflowHostCreds hostCreds) {
+    environment.put("MLFLOW_TRACKING_URI", hostCreds.getHost());
+    if (hostCreds.getUsername() != null) {
+      environment.put("MLFLOW_TRACKING_USERNAME", hostCreds.getUsername());
+    }
+    if (hostCreds.getPassword() != null) {
+      environment.put("MLFLOW_TRACKING_PASSWORD", hostCreds.getPassword());
+    }
+    if (hostCreds.getToken() != null) {
+      environment.put("MLFLOW_TRACKING_TOKEN", hostCreds.getToken());
+    }
+    if (hostCreds.shouldIgnoreTlsVerification()) {
+      environment.put("MLFLOW_TRACKING_INSECURE_TLS", "true");
+    }
+  }
+
+  @VisibleForTesting
+  void setProcessEnvironmentDatabricks(
+      Map<String, String> environment,
+      DatabricksMlflowHostCreds hostCreds) {
+    environment.put("DATABRICKS_HOST", hostCreds.getHost());
+    if (hostCreds.getUsername() != null) {
+      environment.put("DATABRICKS_USERNAME", hostCreds.getUsername());
+    }
+    if (hostCreds.getPassword() != null) {
+      environment.put("DATABRICKS_PASSWORD", hostCreds.getPassword());
+    }
+    if (hostCreds.getToken() != null) {
+      environment.put("DATABRICKS_TOKEN", hostCreds.getToken());
+    }
+    if (hostCreds.shouldIgnoreTlsVerification()) {
+      environment.put("DATABRICKS_INSECURE", "true");
+    }
+  }
+
+  /** Does our best to get the process's stderr, or returns a dummy return value. */
+  private String getErrorBestEffort(Process process) {
+    if (process == null) {
+      return "<process not started>";
+    }
+    try {
+      return IOUtils.toString(process.getErrorStream(), StandardCharsets.UTF_8);
+    } catch (IOException e) {
+      return "<error unknown>";
+    }
+  }
+
+  /** Appends --run-id $runId and --artifact-path $artifactPath, omitting artifactPath if null. */
+  private List<String> appendRunIdArtifactPath(
+      List<String> baseCommand,
+      String runId,
+      String artifactPath) {
+    baseCommand.add("--run-id");
+    baseCommand.add(runId);
+    if (artifactPath != null) {
+      baseCommand.add("--artifact-path");
+      baseCommand.add(artifactPath);
+    }
+    return baseCommand;
+  }
+
+  /** Returns user-facing identifier "runId=abc, artifactId=/foo", omitting artifactPath if null. */
+  private String getTargetIdentifier(String artifactPath) {
+    String identifier = "runId=" + runId;
+    if (artifactPath != null) {
+      return identifier + ", artifactPath=" + artifactPath;
+    }
+    return identifier;
+  }
+}

mlflow/java/client/src/main/java/org/mlflow/internal/proto/Internal.java

@@ -0,0 +1,257 @@
+// Generated by the protocol buffer compiler.  DO NOT EDIT!
+// source: internal.proto
+
+package org.mlflow.internal.proto;
+
+public final class Internal {
+  private Internal() {}
+  public static void registerAllExtensions(
+      com.google.protobuf.ExtensionRegistryLite registry) {
+  }
+
+  public static void registerAllExtensions(
+      com.google.protobuf.ExtensionRegistry registry) {
+    registerAllExtensions(
+        (com.google.protobuf.ExtensionRegistryLite) registry);
+  }
+  /**
+   * <pre>
+   * Types of vertices represented in MLflow Run Inputs. Valid vertices are MLflow objects that can
+   * have an input relationship.
+   * </pre>
+   *
+   * Protobuf enum {@code mlflow.internal.InputVertexType}
+   */
+  public enum InputVertexType
+      implements com.google.protobuf.ProtocolMessageEnum {
+    /**
+     * <code>RUN = 1;</code>
+     */
+    RUN(1),
+    /**
+     * <code>DATASET = 2;</code>
+     */
+    DATASET(2),
+    /**
+     * <code>MODEL = 3;</code>
+     */
+    MODEL(3),
+    ;
+
+    /**
+     * <code>RUN = 1;</code>
+     */
+    public static final int RUN_VALUE = 1;
+    /**
+     * <code>DATASET = 2;</code>
+     */
+    public static final int DATASET_VALUE = 2;
+    /**
+     * <code>MODEL = 3;</code>
+     */
+    public static final int MODEL_VALUE = 3;
+
+
+    public final int getNumber() {
+      return value;
+    }
+
+    /**
+     * @param value The numeric wire value of the corresponding enum entry.
+     * @return The enum associated with the given numeric wire value.
+     * @deprecated Use {@link #forNumber(int)} instead.
+     */
+    @java.lang.Deprecated
+    public static InputVertexType valueOf(int value) {
+      return forNumber(value);
+    }
+
+    /**
+     * @param value The numeric wire value of the corresponding enum entry.
+     * @return The enum associated with the given numeric wire value.
+     */
+    public static InputVertexType forNumber(int value) {
+      switch (value) {
+        case 1: return RUN;
+        case 2: return DATASET;
+        case 3: return MODEL;
+        default: return null;
+      }
+    }
+
+    public static com.google.protobuf.Internal.EnumLiteMap<InputVertexType>
+        internalGetValueMap() {
+      return internalValueMap;
+    }
+    private static final com.google.protobuf.Internal.EnumLiteMap<
+        InputVertexType> internalValueMap =
+          new com.google.protobuf.Internal.EnumLiteMap<InputVertexType>() {
+            public InputVertexType findValueByNumber(int number) {
+              return InputVertexType.forNumber(number);
+            }
+          };
+
+    public final com.google.protobuf.Descriptors.EnumValueDescriptor
+        getValueDescriptor() {
+      return getDescriptor().getValues().get(ordinal());
+    }
+    public final com.google.protobuf.Descriptors.EnumDescriptor
+        getDescriptorForType() {
+      return getDescriptor();
+    }
+    public static final com.google.protobuf.Descriptors.EnumDescriptor
+        getDescriptor() {
+      return org.mlflow.internal.proto.Internal.getDescriptor().getEnumTypes().get(0);
+    }
+
+    private static final InputVertexType[] VALUES = values();
+
+    public static InputVertexType valueOf(
+        com.google.protobuf.Descriptors.EnumValueDescriptor desc) {
+      if (desc.getType() != getDescriptor()) {
+        throw new java.lang.IllegalArgumentException(
+          "EnumValueDescriptor is not for this type.");
+      }
+      return VALUES[desc.getIndex()];
+    }
+
+    private final int value;
+
+    private InputVertexType(int value) {
+      this.value = value;
+    }
+
+    // @@protoc_insertion_point(enum_scope:mlflow.internal.InputVertexType)
+  }
+
+  /**
+   * <pre>
+   * Types of vertices represented in MLflow Run Outputs. Valid vertices are MLflow objects that can
+   * have an output relationship.
+   * </pre>
+   *
+   * Protobuf enum {@code mlflow.internal.OutputVertexType}
+   */
+  public enum OutputVertexType
+      implements com.google.protobuf.ProtocolMessageEnum {
+    /**
+     * <code>RUN_OUTPUT = 1;</code>
+     */
+    RUN_OUTPUT(1),
+    /**
+     * <code>MODEL_OUTPUT = 2;</code>
+     */
+    MODEL_OUTPUT(2),
+    ;
+
+    /**
+     * <code>RUN_OUTPUT = 1;</code>
+     */
+    public static final int RUN_OUTPUT_VALUE = 1;
+    /**
+     * <code>MODEL_OUTPUT = 2;</code>
+     */
+    public static final int MODEL_OUTPUT_VALUE = 2;
+
+
+    public final int getNumber() {
+      return value;
+    }
+
+    /**
+     * @param value The numeric wire value of the corresponding enum entry.
+     * @return The enum associated with the given numeric wire value.
+     * @deprecated Use {@link #forNumber(int)} instead.
+     */
+    @java.lang.Deprecated
+    public static OutputVertexType valueOf(int value) {
+      return forNumber(value);
+    }
+
+    /**
+     * @param value The numeric wire value of the corresponding enum entry.
+     * @return The enum associated with the given numeric wire value.
+     */
+    public static OutputVertexType forNumber(int value) {
+      switch (value) {
+        case 1: return RUN_OUTPUT;
+        case 2: return MODEL_OUTPUT;
+        default: return null;
+      }
+    }
+
+    public static com.google.protobuf.Internal.EnumLiteMap<OutputVertexType>
+        internalGetValueMap() {
+      return internalValueMap;
+    }
+    private static final com.google.protobuf.Internal.EnumLiteMap<
+        OutputVertexType> internalValueMap =
+          new com.google.protobuf.Internal.EnumLiteMap<OutputVertexType>() {
+            public OutputVertexType findValueByNumber(int number) {
+              return OutputVertexType.forNumber(number);
+            }
+          };
+
+    public final com.google.protobuf.Descriptors.EnumValueDescriptor
+        getValueDescriptor() {
+      return getDescriptor().getValues().get(ordinal());
+    }
+    public final com.google.protobuf.Descriptors.EnumDescriptor
+        getDescriptorForType() {
+      return getDescriptor();
+    }
+    public static final com.google.protobuf.Descriptors.EnumDescriptor
+        getDescriptor() {
+      return org.mlflow.internal.proto.Internal.getDescriptor().getEnumTypes().get(1);
+    }
+
+    private static final OutputVertexType[] VALUES = values();
+
+    public static OutputVertexType valueOf(
+        com.google.protobuf.Descriptors.EnumValueDescriptor desc) {
+      if (desc.getType() != getDescriptor()) {
+        throw new java.lang.IllegalArgumentException(
+          "EnumValueDescriptor is not for this type.");
+      }
+      return VALUES[desc.getIndex()];
+    }
+
+    private final int value;
+
+    private OutputVertexType(int value) {
+      this.value = value;
+    }
+
+    // @@protoc_insertion_point(enum_scope:mlflow.internal.OutputVertexType)
+  }
+
+
+  public static com.google.protobuf.Descriptors.FileDescriptor
+      getDescriptor() {
+    return descriptor;
+  }
+  private static  com.google.protobuf.Descriptors.FileDescriptor
+      descriptor;
+  static {
+    java.lang.String[] descriptorData = {
+      "\n\016internal.proto\022\017mlflow.internal\032\025scala" +
+      "pb/scalapb.proto*2\n\017InputVertexType\022\007\n\003R" +
+      "UN\020\001\022\013\n\007DATASET\020\002\022\t\n\005MODEL\020\003*4\n\020OutputVe" +
+      "rtexType\022\016\n\nRUN_OUTPUT\020\001\022\020\n\014MODEL_OUTPUT" +
+      "\020\002B#\n\031org.mlflow.internal.proto\220\001\001\342?\002\020\001"
+    };
+    descriptor = com.google.protobuf.Descriptors.FileDescriptor
+      .internalBuildGeneratedFileFrom(descriptorData,
+        new com.google.protobuf.Descriptors.FileDescriptor[] {
+          org.mlflow.scalapb_interface.Scalapb.getDescriptor(),
+        });
+    com.google.protobuf.ExtensionRegistry registry =
+        com.google.protobuf.ExtensionRegistry.newInstance();
+    registry.add(org.mlflow.scalapb_interface.Scalapb.options);
+    com.google.protobuf.Descriptors.FileDescriptor
+        .internalUpdateFileDescriptor(descriptor, registry);
+    org.mlflow.scalapb_interface.Scalapb.getDescriptor();
+  }
+
+  // @@protoc_insertion_point(outer_class_scope)
+}

mlflow/java/client/src/main/java/org/mlflow/tracking/ActiveRun.java

@@ -0,0 +1,202 @@
+package org.mlflow.tracking;
+
+import org.mlflow.api.proto.Service.*;
+
+import java.nio.file.Path;
+import java.util.*;
+import java.util.stream.Collectors;
+import java.util.stream.StreamSupport;
+
+/**
+ * Represents an active MLflow run and contains APIs to log data to the run.
+ */
+public class ActiveRun {
+  private MlflowClient client;
+  private RunInfo runInfo;
+
+  ActiveRun(RunInfo runInfo, MlflowClient client) {
+    this.runInfo = runInfo;
+    this.client = client;
+  }
+
+  /**
+   * Gets the run id of this run.
+   * @return The run id of this run.
+   */
+  public String getId() {
+    return runInfo.getRunId();
+  }
+
+  /**
+   * Log a parameter under this run.
+   *
+   * @param key The name of the parameter.
+   * @param value The value of the parameter.
+   */
+  public void logParam(String key, String value) {
+    client.logParam(getId(), key, value);
+  }
+
+  /**
+   * Sets a tag under this run.
+   *
+   * @param key The name of the tag.
+   * @param value The value of the tag.
+   */
+  public void setTag(String key, String value) {
+    client.setTag(getId(), key, value);
+  }
+
+  /**
+   * Like {@link #logMetric(String, double, int)} with a default step of 0.
+   */
+  public void logMetric(String key, double value) {
+    logMetric(key, value, 0);
+  }
+
+  /**
+   * Logs a metric under this run.
+   *
+   * @param key The name of the metric.
+   * @param value The value of the metric.
+   * @param step The metric step.
+   */
+  public void logMetric(String key, double value, int step) {
+    client.logMetric(getId(), key, value, System.currentTimeMillis(), step);
+  }
+
+  /**
+   * Like {@link #logMetrics(Map, int)} with a default step of 0.
+   */
+  public void logMetrics(Map<String, Double> metrics) {
+      logMetrics(metrics, 0);
+  }
+
+  /**
+   * Log multiple metrics for this run.
+   *
+   * @param metrics A map of metric name to value.
+   * @param step The metric step.
+   */
+  public void logMetrics(Map<String, Double> metrics, int step) {
+    List<Metric> protoMetrics = metrics.entrySet().stream()
+      .map((metric) ->
+        Metric.newBuilder()
+          .setKey(metric.getKey())
+          .setValue(metric.getValue())
+          .setTimestamp(System.currentTimeMillis())
+          .setStep(step)
+          .build()
+      ).collect(Collectors.toList());
+    client.logBatch(getId(), protoMetrics, Collections.emptyList(), Collections.emptyList());
+  }
+
+  /**
+   * Log multiple params for this run.
+   *
+   * @param params A map of param name to value.
+   */
+  public void logParams(Map<String, String> params) {
+    List<Param> protoParams = params.entrySet().stream().map((param) ->
+      Param.newBuilder()
+        .setKey(param.getKey())
+        .setValue(param.getValue())
+        .build()
+    ).collect(Collectors.toList());
+    client.logBatch(getId(), Collections.emptyList(), protoParams, Collections.emptyList());
+  }
+
+  /**
+   * Sets multiple tags for this run.
+   *
+   * @param tags A map of tag name to value.
+   */
+  public void setTags(Map<String, String> tags) {
+    List<RunTag> protoTags = tags.entrySet().stream().map((tag) ->
+      RunTag.newBuilder().setKey(tag.getKey()).setValue(tag.getValue()).build()
+    ).collect(Collectors.toList());
+    client.logBatch(getId(), Collections.emptyList(), Collections.emptyList(), protoTags);
+  }
+
+  /**
+   * Like {@link #logArtifact(Path, String)} with the artifactPath set to the root of the
+   * artifact directory.
+   *
+   * @param localPath Path of file to upload. Must exist, and must be a simple file
+   *                  (not a directory).
+   */
+  public void logArtifact(Path localPath) {
+    client.logArtifact(getId(), localPath.toFile());
+  }
+
+  /**
+   * Uploads the given local file to the run's root artifact directory. For example,
+   *
+   *   <pre>
+   *   activeRun.logArtifact("/my/localModel", "model")
+   *   mlflowClient.listArtifacts(activeRun.getId(), "model") // returns "model/localModel"
+   *   </pre>
+   *
+   * @param localPath Path of file to upload. Must exist, and must be a simple file
+   *                  (not a directory).
+   * @param artifactPath Artifact path relative to the run's root directory given by
+   *                     {@link #getArtifactUri()}. Should NOT start with a /.
+   */
+  public void logArtifact(Path localPath, String artifactPath) {
+    client.logArtifact(getId(), localPath.toFile(), artifactPath);
+  }
+
+  /**
+   * Like {@link #logArtifacts(Path, String)} with the artifactPath set to the root of the
+   * artifact directory.
+   *
+   * @param localPath Directory to upload. Must exist, and must be a directory (not a simple file).
+   */
+  public void logArtifacts(Path localPath) {
+    client.logArtifacts(getId(), localPath.toFile());
+  }
+
+  /**
+   * Uploads all files within the given local director an artifactPath within the run's root
+   * artifact directory. For example, if /my/local/dir/ contains two files "file1" and "file2", then
+   *
+   *   <pre>
+   *   activeRun.logArtifacts("/my/local/dir", "model")
+   *   mlflowClient.listArtifacts(activeRun.getId(), "model") // returns "model/file1" and
+   *                                                          // "model/file2"
+   *   </pre>
+   *
+   * (i.e., the contents of the local directory are now available in model/).
+   *
+   * @param localPath Directory to upload. Must exist, and must be a directory (not a simple file).
+   * @param artifactPath Artifact path relative to the run's root directory given by
+   *                     {@link #getArtifactUri()}. Should NOT start with a /.
+   */
+  public void logArtifacts(Path localPath, String artifactPath) {
+    client.logArtifacts(getId(), localPath.toFile(), artifactPath);
+  }
+
+  /**
+   * Get the absolute URI of the run artifact directory root.
+   * @return The absolute URI of the run artifact directory root.
+   */
+  public String getArtifactUri() {
+    return this.runInfo.getArtifactUri();
+  }
+
+  /**
+   * Ends the active MLflow run.
+   */
+  public void endRun() {
+    endRun(RunStatus.FINISHED);
+  }
+
+  /**
+   * Ends the active MLflow run.
+   *
+   * @param status The status of the run.
+   */
+  public void endRun(RunStatus status) {
+    client.setTerminated(getId(), status);
+  }
+}

mlflow/java/client/src/main/java/org/mlflow/tracking/EmptyPage.java

@@ -0,0 +1,48 @@
+package org.mlflow.tracking;
+
+import java.util.Collections;
+import java.util.Optional;
+
+public class EmptyPage<E> implements Page<E> {
+
+  /**
+   * Creates an empty page
+   */
+  EmptyPage() {}
+
+  /**
+   * @return Zero
+   */
+  public int getPageSize() {
+    return 0;
+  }
+
+  /**
+   * @return False
+   */
+  public boolean hasNextPage() {
+    return false;
+  }
+
+  /**
+   * @return An empty Optional.
+   */
+  public Optional<String> getNextPageToken() {
+    return Optional.empty();
+  }
+
+  /**
+   * @return An {@link org.mlflow.tracking.EmptyPage}
+   */
+  public EmptyPage getNextPage() {
+    return this;
+  }
+
+  /**
+   * @return An empty iterable.
+   */
+  public Iterable getItems() {
+    return Collections.EMPTY_LIST;
+  }
+
+}

mlflow/java/client/src/main/java/org/mlflow/tracking/ExperimentsPage.java

@@ -0,0 +1,88 @@
+package org.mlflow.tracking;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.ArrayList;
+import java.util.Optional;
+import org.mlflow.api.proto.Service.*;
+
+public class ExperimentsPage implements Page<Experiment> {
+
+  private final String token;
+  private final List<Experiment> experiments;
+
+  private final MlflowClient client;
+  private final String searchFilter;
+  private final ViewType experimentViewType;
+  private final List<String> orderBy;
+  private final int maxResults;
+
+  /**
+   * Creates a fixed size page of Experiments.
+   */
+  ExperimentsPage(List<Experiment> experiments,
+                  String token,
+                  String searchFilter,
+                  ViewType experimentViewType,
+                  int maxResults,
+                  List<String> orderBy,
+                  MlflowClient client) {
+    this.experiments = Collections.unmodifiableList(experiments);
+    this.token = token;
+    this.searchFilter = searchFilter;
+    this.experimentViewType = experimentViewType;
+    this.orderBy = orderBy;
+    this.maxResults = maxResults;
+    this.client = client;
+  }
+
+  /**
+   * @return The number of experiments in the page.
+   */
+  public int getPageSize() {
+    return this.experiments.size();
+  }
+
+  /**
+   * @return True if a token for the next page exists and isn't empty. Otherwise returns false.
+   */
+  public boolean hasNextPage() {
+    return this.token != null && this.token != "";
+  }
+
+  /**
+   * @return An optional with the token for the next page. 
+   * Empty if the token doesn't exist or is empty.
+   */
+  public Optional<String> getNextPageToken() {
+    if (this.hasNextPage()) {
+      return Optional.of(this.token);
+    } else {
+      return Optional.empty();
+    }
+  }
+
+  /**
+   * @return The next page of experiments matching the search criteria.
+   * If there are no more pages, an {@link org.mlflow.tracking.EmptyPage} will be returned.
+   */
+  public Page<Experiment> getNextPage() {
+    if (this.hasNextPage()) {
+      return this.client.searchExperiments(this.searchFilter,
+                                           this.experimentViewType,
+                                           this.maxResults,
+                                           this.orderBy,
+                                           this.token);
+    } else {
+      return new EmptyPage();
+    }
+  }
+
+  /**
+   * @return An iterable over the experiments in this page.
+   */
+  public List<Experiment> getItems() {
+    return experiments;
+  }
+
+}

mlflow/java/client/src/main/java/org/mlflow/tracking/MlflowClientException.java

@@ -0,0 +1,14 @@
+package org.mlflow.tracking;
+
+/** Superclass of all exceptions thrown by the MlflowClient API. */
+public class MlflowClientException extends RuntimeException {
+  public MlflowClientException(String message) {
+    super(message);
+  }
+  public MlflowClientException(String message, Throwable cause) {
+    super(message, cause);
+  }
+  public MlflowClientException(Throwable cause) {
+    super(cause);
+  }
+}

mlflow/java/client/src/main/java/org/mlflow/tracking/MlflowClientVersion.java

@@ -0,0 +1,34 @@
+package org.mlflow.tracking;
+
+import java.io.InputStream;
+import java.util.Properties;
+
+import com.google.common.base.Supplier;
+import com.google.common.base.Suppliers;
+
+/** Returns the version of the MLflow project this client was compiled against. */
+public class MlflowClientVersion {
+  // To avoid extra disk IO during class loading (static initialization), we lazily read the
+  // pom.properties file on first access and then cache the result to avoid future IO.
+  private static Supplier<String> clientVersionSupplier = Suppliers.memoize(() -> {
+    try {
+      Properties p = new Properties();
+      InputStream is = MlflowClientVersion.class.getResourceAsStream(
+        "/META-INF/maven/org.mlflow/mlflow-client/pom.properties");
+      if (is == null) {
+        return "";
+      }
+      p.load(is);
+      return p.getProperty("version", "");
+    } catch (Exception e) {
+      return "";
+    }
+  });
+
+  private MlflowClientVersion() {}
+
+  /** @return MLflow client version (e.g., 0.9.1) or an empty string if detection fails. */
+  public static String getClientVersion() {
+    return clientVersionSupplier.get();
+  }
+}

mlflow/java/client/src/main/java/org/mlflow/tracking/MlflowContext.java

@@ -0,0 +1,255 @@
+package org.mlflow.tracking;
+
+import org.mlflow.api.proto.Service.*;
+import org.mlflow.tracking.utils.DatabricksContext;
+import org.mlflow.tracking.utils.MlflowTagConstants;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.*;
+import java.util.function.Consumer;
+
+/**
+ * Main entrypoint used to start MLflow runs to log to. This is a higher level interface than
+ * {@code MlflowClient} and provides convenience methods to keep track of active runs and to set
+ * default tags on runs which are created through {@code MlflowContext}
+ *
+ * On construction, MlflowContext will choose a default experiment ID to log to depending on your
+ * environment. To log to a different experiment, use {@link #setExperimentId(String)} or
+ * {@link #setExperimentName(String)}
+ *
+ * <p>
+ * For example:
+ *   <pre>
+ *   // Uses the URI set in the MLFLOW_TRACKING_URI environment variable.
+ *   // To use your own tracking uri set it in the call to "new MlflowContext("tracking-uri")"
+ *   MlflowContext mlflow = new MlflowContext();
+ *   ActiveRun run = mlflow.startRun("run-name");
+ *   run.logParam("alpha", "0.5");
+ *   run.logMetric("MSE", 0.0);
+ *   run.endRun();
+ *   </pre>
+ */
+public class MlflowContext {
+  private MlflowClient client;
+  private String experimentId;
+  // Cache the default experiment ID for a repo notebook to avoid sending
+  // extraneous API requests
+  private static String defaultRepoNotebookExperimentId;
+  private static final Logger logger = LoggerFactory.getLogger(MlflowContext.class);
+
+
+  /**
+   * Constructs a {@code MlflowContext} with a MlflowClient based on the MLFLOW_TRACKING_URI
+   * environment variable.
+   */
+  public MlflowContext() {
+    this(new MlflowClient());
+  }
+
+  /**
+   * Constructs a {@code MlflowContext} which points to the specified trackingUri.
+   *
+   * @param trackingUri The URI to log to.
+   */
+  public MlflowContext(String trackingUri) {
+    this(new MlflowClient(trackingUri));
+  }
+
+  /**
+   * Constructs a {@code MlflowContext} which points to the specified trackingUri.
+   *
+   * @param client The client used to log runs.
+   */
+  public MlflowContext(MlflowClient client) {
+    this.client = client;
+    this.experimentId = getDefaultExperimentId();
+  }
+
+  /**
+   * Returns the client used to log runs.
+   *
+   * @return the client used to log runs.
+   */
+  public MlflowClient getClient() {
+    return this.client;
+  }
+
+  /**
+   * Sets the experiment to log runs to by name.
+   * @param experimentName the name of the experiment to log runs to.
+   * @throws IllegalArgumentException if the experiment name does not match an existing experiment
+   */
+  public MlflowContext setExperimentName(String experimentName) throws IllegalArgumentException {
+    Optional<Experiment> experimentOpt = client.getExperimentByName(experimentName);
+    if (!experimentOpt.isPresent()) {
+      throw new IllegalArgumentException(
+        String.format("%s is not a valid experiment", experimentName));
+    }
+    experimentId = experimentOpt.get().getExperimentId();
+    return this;
+  }
+
+  /**
+   * Sets the experiment to log runs to by ID.
+   * @param experimentId the id of the experiment to log runs to.
+   */
+  public MlflowContext setExperimentId(String experimentId) {
+    this.experimentId = experimentId;
+    return this;
+  }
+
+  /**
+   * Returns the experiment ID we are logging to.
+   *
+   * @return the experiment ID we are logging to.
+   */
+  public String getExperimentId() {
+    return this.experimentId;
+  }
+
+  /**
+   * Starts a MLflow run without a name. To log data to newly created MLflow run see the methods on
+   * {@link ActiveRun}. MLflow runs should be ended using {@link ActiveRun#endRun()}
+   *
+   * @return An {@code ActiveRun} object to log data to.
+   */
+  public ActiveRun startRun() {
+    return startRun(null);
+  }
+
+  /**
+   * Starts a MLflow run. To log data to newly created MLflow run see the methods on
+   * {@link ActiveRun}. MLflow runs should be ended using {@link ActiveRun#endRun()}
+   *
+   * @param runName The name of this run. For display purposes only and is stored in the
+   *                mlflow.runName tag.
+   * @return An {@code ActiveRun} object to log data to.
+   */
+  public ActiveRun startRun(String runName) {
+    return startRun(runName, null);
+  }
+
+  /**
+   * Like {@link #startRun(String)} but sets the {@code mlflow.parentRunId} tag in order to create
+   * nested runs.
+   *
+   * @param runName The name of this run. For display purposes only and is stored in the
+   *                mlflow.runName tag.
+   * @param parentRunId The ID of this run's parent
+   * @return An {@code ActiveRun} object to log data to.
+   */
+  public ActiveRun startRun(String runName, String parentRunId) {
+    Map<String, String> tags = new HashMap<>();
+    if (runName != null) {
+      tags.put(MlflowTagConstants.RUN_NAME, runName);
+    }
+    tags.put(MlflowTagConstants.USER, System.getProperty("user.name"));
+    tags.put(MlflowTagConstants.SOURCE_TYPE, "LOCAL");
+    if (parentRunId != null) {
+      tags.put(MlflowTagConstants.PARENT_RUN_ID, parentRunId);
+    }
+
+    // Add tags from DatabricksContext if they exist
+    DatabricksContext databricksContext = DatabricksContext.createIfAvailable();
+    if (databricksContext != null) {
+      tags.putAll(databricksContext.getTags());
+    }
+
+    CreateRun.Builder createRunBuilder = CreateRun.newBuilder()
+      .setExperimentId(experimentId)
+      .setStartTime(System.currentTimeMillis());
+    for (Map.Entry<String, String> tag: tags.entrySet()) {
+      createRunBuilder.addTags(
+        RunTag.newBuilder().setKey(tag.getKey()).setValue(tag.getValue()).build());
+    }
+    RunInfo runInfo = client.createRun(createRunBuilder.build());
+
+    ActiveRun newRun = new ActiveRun(runInfo, client);
+    return newRun;
+  }
+
+  /**
+   * Like {@link #startRun(String)} but will terminate the run after the activeRunFunction is
+   * executed.
+   *
+   * For example
+   *   <pre>
+   *   mlflowContext.withActiveRun((activeRun -&gt; {
+   *     activeRun.logParam("layers", "4");
+   *   }));
+   *   </pre>
+   *
+   * @param activeRunFunction A function which takes an {@code ActiveRun} and logs data to it.
+   */
+  public void withActiveRun(Consumer<ActiveRun> activeRunFunction) {
+    ActiveRun newRun = startRun();
+    try {
+      activeRunFunction.accept(newRun);
+    } catch(Exception e) {
+      newRun.endRun(RunStatus.FAILED);
+      return;
+    }
+    newRun.endRun(RunStatus.FINISHED);
+  }
+
+  /**
+   * Like {@link #withActiveRun(Consumer)} with an explicity run name.
+   *
+   * @param runName The name of this run. For display purposes only and is stored in the
+   *                mlflow.runName tag.
+   * @param activeRunFunction A function which takes an {@code ActiveRun} and logs data to it.
+   */
+  public void withActiveRun(String runName, Consumer<ActiveRun> activeRunFunction) {
+    ActiveRun newRun = startRun(runName);
+    try {
+      activeRunFunction.accept(newRun);
+    } catch(Exception e) {
+      newRun.endRun(RunStatus.FAILED);
+      return;
+    }
+    newRun.endRun(RunStatus.FINISHED);
+  }
+
+  private static String getDefaultRepoNotebookExperimentId(String notebookId, String notebookPath) {
+      if (defaultRepoNotebookExperimentId != null) {
+        return defaultRepoNotebookExperimentId;
+      }
+      CreateExperiment.Builder request = CreateExperiment.newBuilder();
+      request.setName(notebookPath);
+      request.addTags(ExperimentTag.newBuilder()
+              .setKey(MlflowTagConstants.MLFLOW_EXPERIMENT_SOURCE_TYPE)
+              .setValue("REPO_NOTEBOOK")
+      );
+      request.addTags(ExperimentTag.newBuilder()
+              .setKey(MlflowTagConstants.MLFLOW_EXPERIMENT_SOURCE_ID)
+              .setValue(notebookId)
+      );
+      String experimentId = (new MlflowClient()).createExperiment(request.build());
+      defaultRepoNotebookExperimentId = experimentId;
+      return experimentId;
+
+  }
+
+  private static String getDefaultExperimentId() {
+    DatabricksContext databricksContext = DatabricksContext.createIfAvailable();
+    if (databricksContext != null && databricksContext.isInDatabricksNotebook()) {
+      String notebookId = databricksContext.getNotebookId();
+      String notebookPath = databricksContext.getNotebookPath();
+      if (notebookId != null) {
+        if (notebookPath != null && notebookPath.startsWith("/Repos")) {
+          try {
+            return getDefaultRepoNotebookExperimentId(notebookId, notebookPath);
+          }
+          catch (Exception e) {
+            // Do nothing; will fall through to returning notebookId
+            logger.warn("Failed to get default repo notebook experiment ID", e);
+          }
+        }
+        return notebookId;
+      }
+    }
+    return MlflowClient.DEFAULT_EXPERIMENT_ID;
+  }
+}

mlflow/java/client/src/main/java/org/mlflow/tracking/MlflowHttpCaller.java

@@ -0,0 +1,260 @@
+package org.mlflow.tracking;
+
+import com.google.common.annotations.VisibleForTesting;
+
+import java.io.IOException;
+import java.net.URI;
+import java.nio.charset.StandardCharsets;
+import java.security.KeyManagementException;
+import java.security.KeyStoreException;
+import java.security.NoSuchAlgorithmException;
+import java.util.Base64;
+
+import org.apache.http.HttpResponse;
+import org.apache.http.client.methods.HttpEntityEnclosingRequestBase;
+import org.apache.http.client.methods.HttpGet;
+import org.apache.http.client.methods.HttpPatch;
+import org.apache.http.client.methods.HttpPost;
+import org.apache.http.client.methods.HttpRequestBase;
+import org.apache.http.conn.ssl.NoopHostnameVerifier;
+import org.apache.http.conn.ssl.SSLConnectionSocketFactory;
+import org.apache.http.conn.ssl.TrustSelfSignedStrategy;
+import org.apache.http.entity.StringEntity;
+import org.apache.http.impl.client.CloseableHttpClient;
+import org.apache.http.impl.client.HttpClientBuilder;
+import org.apache.http.ssl.SSLContextBuilder;
+import org.apache.http.util.EntityUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.mlflow.tracking.creds.MlflowHostCreds;
+import org.mlflow.tracking.creds.MlflowHostCredsProvider;
+
+
+class MlflowHttpCaller {
+  private static final Logger logger = LoggerFactory.getLogger(MlflowHttpCaller.class);
+  private static final String BASE_API_PATH = "api/2.0/mlflow";
+  protected CloseableHttpClient httpClient;
+  private final MlflowHostCredsProvider hostCredsProvider;
+  private final int maxRateLimitIntervalMillis;
+  private final int rateLimitRetrySleepInitMillis;
+  private final int maxRetryAttempts;
+
+  /**
+   * Construct a new MlflowHttpCaller with a default configuration for request retries.
+   */
+  MlflowHttpCaller(MlflowHostCredsProvider hostCredsProvider) {
+    this(hostCredsProvider, 60000, 1000, 3);
+  }
+
+  /**
+   * Construct a new MlflowHttpCaller.
+   *
+   * @param maxRateLimitIntervalMs The maximum amount of time, in milliseconds, to spend retrying a
+   *                               single request in response to rate limiting (error code 429).
+   * @param rateLimitRetrySleepInitMs The initial backoff delay, in milliseconds, when retrying a
+   *                                  request in response to rate limiting (error code 429). The
+   *                                  delay is increased exponentially after each rate limiting
+   *                                  response until the total delay incurred across all retries for
+   *                                  the request exceeds the specified maxRateLimitIntervalSeconds.
+   * @param maxRetryAttempts The maximum number of times to retry a request, excluding rate limit
+   *                         retries.
+   */
+  MlflowHttpCaller(MlflowHostCredsProvider hostCredsProvider,
+                   int maxRateLimitIntervalMs,
+                   int rateLimitRetrySleepInitMs,
+                   int maxRetryAttempts) {
+    this.hostCredsProvider = hostCredsProvider;
+    this.maxRateLimitIntervalMillis = maxRateLimitIntervalMs;
+    this.rateLimitRetrySleepInitMillis = rateLimitRetrySleepInitMs;
+    this.maxRetryAttempts = maxRetryAttempts;
+  }
+
+  @VisibleForTesting
+  MlflowHttpCaller(MlflowHostCredsProvider hostCredsProvider,
+                   int maxRateLimitIntervalMs,
+                   int rateLimitRetrySleepInitMs,
+                   int maxRetryAttempts,
+                   CloseableHttpClient client) {
+    this(
+      hostCredsProvider, maxRateLimitIntervalMs, rateLimitRetrySleepInitMs, maxRetryAttempts);
+    this.httpClient = client;
+  }
+
+  private HttpResponse executeRequestWithRateLimitRetries(HttpRequestBase request)
+      throws IOException {
+    int timeLeft = maxRateLimitIntervalMillis;
+    int sleepFor = rateLimitRetrySleepInitMillis;
+    HttpResponse response = httpClient.execute(request);
+    while (response.getStatusLine().getStatusCode() == 429 && timeLeft > 0) {
+      logger.warn("Request returned with status code 429 (Rate limit exceeded). Retrying after "
+                  + sleepFor
+                  + " milliseconds. Will continue to retry 429s for up to "
+                  + timeLeft
+                  + " milliseconds.");
+      try {
+        Thread.sleep(sleepFor);
+      } catch (InterruptedException e) {
+        throw new RuntimeException(e);
+      }
+      timeLeft -= sleepFor;
+      sleepFor = Math.min(timeLeft, 2 * sleepFor);
+      response = httpClient.execute(request);
+    }
+    checkError(response);
+    return response;
+  }
+
+  private HttpResponse executeRequest(HttpRequestBase request) throws IOException {
+    HttpResponse response = null;
+    int attemptsRemaining = this.maxRetryAttempts;
+    while (attemptsRemaining > 0) {
+      attemptsRemaining -= 1;
+      try {
+        response = executeRequestWithRateLimitRetries(request);
+        break;
+      } catch (MlflowHttpException e) {
+        if (attemptsRemaining > 0 && e.getStatusCode() != 429) {
+          logger.warn("Request returned with status code {} (Rate limit exceeded)."
+                      + " Retrying up to {} more times. Response body: {}",
+                      e.getStatusCode(),
+                      attemptsRemaining,
+                      e.getBodyMessage());
+          continue;
+        } else {
+          throw e;
+        }
+      }
+    }
+    return response;
+  }
+
+  String get(String path) {
+    logger.debug("Sending GET " + path);
+    HttpGet request = new HttpGet();
+    fillRequestSettings(request, path);
+    try {
+      HttpResponse response = executeRequest(request);
+      String responseJson = EntityUtils.toString(response.getEntity(), StandardCharsets.UTF_8);
+      logger.debug("Response: " + responseJson);
+      return responseJson;
+    } catch (IOException e) {
+      throw new MlflowClientException(e);
+    }
+  }
+
+  // TODO(aaron) Convert to InputStream.
+  byte[] getAsBytes(String path) {
+    logger.debug("Sending GET " + path);
+    HttpGet request = new HttpGet();
+    fillRequestSettings(request, path);
+    try {
+      HttpResponse response = executeRequest(request);
+      byte[] bytes = EntityUtils.toByteArray(response.getEntity());
+      logger.debug("response: #bytes=" + bytes.length);
+      return bytes;
+    } catch (IOException e) {
+      throw new MlflowClientException(e);
+    }
+  }
+
+  String post(String path, String json) {
+    logger.debug("Sending POST " + path + ": " + json);
+    HttpPost request = new HttpPost();
+    return send(request, path, json);
+  }
+
+  String patch(String path, String json) {
+    logger.debug("Sending PATCH " + path + ": " + json);
+    HttpPatch request = new HttpPatch();
+    return send(request, path, json);
+  }
+
+  private String send(HttpEntityEnclosingRequestBase request, String path, String json) {
+    fillRequestSettings(request, path);
+    request.setEntity(new StringEntity(json, StandardCharsets.UTF_8));
+    request.setHeader("Content-Type", "application/json");
+    try {
+      HttpResponse response = executeRequest(request);
+      String responseJson = EntityUtils.toString(response.getEntity(), StandardCharsets.UTF_8);
+      logger.debug("Response: " + responseJson);
+      return responseJson;
+    } catch (IOException e) {
+      throw new MlflowClientException(e);
+    }
+  }
+
+  private void checkError(HttpResponse response) throws MlflowClientException, IOException {
+    int statusCode = response.getStatusLine().getStatusCode();
+    String reasonPhrase = response.getStatusLine().getReasonPhrase();
+    if (isError(statusCode)) {
+      String bodyMessage = EntityUtils.toString(response.getEntity(), StandardCharsets.UTF_8);
+      if (statusCode >= 400 && statusCode <= 499) {
+        throw new MlflowHttpException(statusCode, reasonPhrase, bodyMessage);
+      }
+      if (statusCode >= 500 && statusCode <= 599) {
+        throw new MlflowHttpException(statusCode, reasonPhrase, bodyMessage);
+      }
+      throw new MlflowHttpException(statusCode, reasonPhrase, bodyMessage);
+    }
+  }
+
+  private void fillRequestSettings(HttpRequestBase request, String path) {
+    MlflowHostCreds hostCreds = hostCredsProvider.getHostCreds();
+    createHttpClientIfNecessary(hostCreds.shouldIgnoreTlsVerification());
+    String uri = hostCreds.getHost() + "/" + BASE_API_PATH + "/" + path;
+    request.setURI(URI.create(uri));
+    String username = hostCreds.getUsername();
+    String password = hostCreds.getPassword();
+    String token = hostCreds.getToken();
+    if (username != null && password != null) {
+      String authHeader = Base64.getEncoder()
+        .encodeToString((username + ":" + password).getBytes(StandardCharsets.UTF_8));
+      request.addHeader("Authorization", "Basic " + authHeader);
+    } else if (token != null) {
+      request.addHeader("Authorization", "Bearer " + token);
+    }
+
+    String userAgent = "mlflow-java-client";
+    String clientVersion = MlflowClientVersion.getClientVersion();
+    if (!clientVersion.isEmpty()) {
+      userAgent += "/" + clientVersion;
+    }
+    request.addHeader("User-Agent", userAgent);
+  }
+
+  private boolean isError(int statusCode) {
+    return statusCode < 200 || statusCode > 399;
+  }
+
+  private void createHttpClientIfNecessary(boolean noTlsVerify) {
+    if (httpClient != null) {
+      return;
+    }
+
+    HttpClientBuilder builder = HttpClientBuilder.create();
+    if (noTlsVerify) {
+      try {
+        SSLContextBuilder sslBuilder = new SSLContextBuilder()
+          .loadTrustMaterial(null, new TrustSelfSignedStrategy());
+        SSLConnectionSocketFactory connectionFactory =
+          new SSLConnectionSocketFactory(sslBuilder.build(), new NoopHostnameVerifier());
+        builder.setSSLSocketFactory(connectionFactory);
+      } catch (NoSuchAlgorithmException | KeyStoreException | KeyManagementException e) {
+        logger.warn("Could not set noTlsVerify to true, verification will remain", e);
+      }
+    }
+
+    this.httpClient = builder.build();
+  }
+
+  void close() {
+    if (httpClient != null) {
+      try {
+	  httpClient.close();
+      } catch(IOException e){
+	  logger.warn("Unable to close connection to mlflow backend", e);
+      }
+    }
+  }
+}

mlflow/java/client/src/main/java/org/mlflow/tracking/MlflowHttpException.java

@@ -0,0 +1,39 @@
+package org.mlflow.tracking;
+
+/**
+ * Returned when an HTTP API request to a remote Tracking service returns an error code.
+ */
+public class MlflowHttpException extends MlflowClientException {
+
+  public MlflowHttpException(int statusCode, String reasonPhrase) {
+    super("statusCode=" + statusCode + " reasonPhrase=[" + reasonPhrase +"]");
+    this.statusCode = statusCode;
+    this.reasonPhrase = reasonPhrase;
+  }
+
+  public MlflowHttpException(int statusCode, String reasonPhrase, String bodyMessage) {
+    super("statusCode=" + statusCode + " reasonPhrase=[" + reasonPhrase + "] bodyMessage=["
+      + bodyMessage + "]");
+    this.statusCode = statusCode;
+    this.reasonPhrase = reasonPhrase;
+    this.bodyMessage = bodyMessage;
+  }
+
+  private int statusCode;
+
+  public int getStatusCode() {
+    return statusCode;
+  }
+
+  private String reasonPhrase;
+
+  public String getReasonPhrase() {
+    return reasonPhrase;
+  }
+
+  private String bodyMessage;
+
+  public String getBodyMessage() {
+    return bodyMessage;
+  }
+}

mlflow/java/client/src/main/java/org/mlflow/tracking/MlflowProtobufMapper.java

@@ -0,0 +1,321 @@
+package org.mlflow.tracking;
+
+import com.google.protobuf.InvalidProtocolBufferException;
+import com.google.protobuf.MessageOrBuilder;
+import com.google.protobuf.util.JsonFormat;
+
+import java.lang.Iterable;
+import java.net.URISyntaxException;
+import java.util.List;
+
+import org.apache.http.client.utils.URIBuilder;
+import org.mlflow.api.proto.ModelRegistry.*;
+import org.mlflow.api.proto.Service.*;
+
+class MlflowProtobufMapper {
+
+  String makeCreateExperimentRequest(String expName) {
+    CreateExperiment.Builder builder = CreateExperiment.newBuilder();
+    builder.setName(expName);
+    return print(builder);
+  }
+
+  String makeDeleteExperimentRequest(String experimentId) {
+    DeleteExperiment.Builder builder = DeleteExperiment.newBuilder();
+    builder.setExperimentId(experimentId);
+    return print(builder);
+  }
+
+  String makeRestoreExperimentRequest(String experimentId) {
+    RestoreExperiment.Builder builder = RestoreExperiment.newBuilder();
+    builder.setExperimentId(experimentId);
+    return print(builder);
+  }
+
+  String makeUpdateExperimentRequest(String experimentId, String newExperimentName) {
+    UpdateExperiment.Builder builder = UpdateExperiment.newBuilder();
+    builder.setExperimentId(experimentId);
+    builder.setNewName(newExperimentName);
+    return print(builder);
+  }
+
+  String makeLogParam(String runId, String key, String value) {
+    LogParam.Builder builder = LogParam.newBuilder();
+    builder.setRunUuid(runId);
+    builder.setRunId(runId);
+    builder.setKey(key);
+    builder.setValue(value);
+    return print(builder);
+  }
+
+  String makeLogMetric(String runId, String key, double value, long timestamp, long step) {
+    LogMetric.Builder builder = LogMetric.newBuilder();
+    builder.setRunUuid(runId);
+    builder.setRunId(runId);
+    builder.setKey(key);
+    builder.setValue(value);
+    builder.setTimestamp(timestamp);
+    builder.setStep(step);
+    return print(builder);
+  }
+
+  String makeSetExperimentTag(String expId, String key, String value) {
+    SetExperimentTag.Builder builder = SetExperimentTag.newBuilder();
+    builder.setExperimentId(expId);
+    builder.setKey(key);
+    builder.setValue(value);
+    return print(builder);
+  }
+
+  String makeSetTag(String runId, String key, String value) {
+    SetTag.Builder builder = SetTag.newBuilder();
+    builder.setRunUuid(runId);
+    builder.setRunId(runId);
+    builder.setKey(key);
+    builder.setValue(value);
+    return print(builder);
+  }
+
+  String makeDeleteTag(String runId, String key) {
+    DeleteTag.Builder builder = DeleteTag.newBuilder();
+    builder.setRunId(runId);
+    builder.setKey(key);
+    return print(builder);
+  }
+
+  String makeLogBatch(String runId,
+      Iterable<Metric> metrics,
+      Iterable<Param> params,
+      Iterable<RunTag> tags) {
+    LogBatch.Builder builder = LogBatch.newBuilder();
+    builder.setRunId(runId);
+    if (metrics != null) {
+      builder.addAllMetrics(metrics);
+    }
+    if (params != null) {
+      builder.addAllParams(params);
+    }
+    if (tags != null) {
+      builder.addAllTags(tags);
+    }
+    return print(builder);
+  }
+
+  String makeUpdateRun(String runId, RunStatus status, long endTime) {
+    UpdateRun.Builder builder = UpdateRun.newBuilder();
+    builder.setRunUuid(runId);
+    builder.setRunId(runId);
+    builder.setStatus(status);
+    builder.setEndTime(endTime);
+    return print(builder);
+  }
+
+  String makeDeleteRun(String runId) {
+    DeleteRun.Builder builder = DeleteRun.newBuilder();
+    builder.setRunId(runId);
+    return print(builder);
+  }
+
+  String makeRestoreRun(String runId) {
+    RestoreRun.Builder builder = RestoreRun.newBuilder();
+    builder.setRunId(runId);
+    return print(builder);
+  }
+
+  String makeGetLatestVersion(String modelName, Iterable<String> stages) {
+    try {
+      URIBuilder builder = new URIBuilder("registered-models/get-latest-versions")
+              .addParameter("name", modelName);
+      if (stages != null) {
+        for( String stage: stages) {
+          builder.addParameter("stages", stage);
+        }
+      }
+      return builder.build().toString();
+    } catch (URISyntaxException e) {
+      throw new MlflowClientException("Failed to construct request URI for get latest versions.",
+              e);
+    }
+  }
+
+  String makeUpdateModelVersion(String modelName, String version) {
+    return print(UpdateModelVersion.newBuilder().setName(modelName).setVersion(version));
+  }
+
+  String makeTransitionModelVersionStage(String modelName, String version, String stage) {
+    return print(TransitionModelVersionStage.newBuilder()
+            .setName(modelName).setVersion(version).setStage(stage));
+  }
+
+  String makeCreateModel(String modelName) {
+    CreateRegisteredModel.Builder builder = CreateRegisteredModel.newBuilder()
+            .setName(modelName);
+    return print(builder);
+  }
+
+  String makeCreateModelVersion(String modelName, String runId, String source) {
+    CreateModelVersion.Builder builder = CreateModelVersion.newBuilder()
+            .setName(modelName)
+            .setRunId(runId)
+            .setSource(source);
+    return print(builder);
+  }
+
+  String makeGetRegisteredModel(String modelName) {
+    try {
+      return new URIBuilder("registered-models/get")
+          .addParameter("name", modelName)
+          .build()
+          .toString();
+    } catch (URISyntaxException e) {
+      throw new MlflowClientException("Failed to construct request URI for get model version.", e);
+    }
+  }
+
+  String makeGetModelVersion(String modelName, String modelVersion) {
+    try {
+      return new URIBuilder("model-versions/get")
+          .addParameter("name", modelName)
+          .addParameter("version", modelVersion)
+          .build()
+          .toString();
+    } catch (URISyntaxException e) {
+      throw new MlflowClientException("Failed to construct request URI for get model version.", e);
+    }
+  }
+
+  String makeGetModelVersionDownloadUri(String modelName, String modelVersion) {
+    try {
+      return new URIBuilder("model-versions/get-download-uri")
+              .addParameter("name", modelName)
+              .addParameter("version", modelVersion).build().toString();
+    } catch (URISyntaxException e) {
+      throw new MlflowClientException(
+              "Failed to construct request URI for get version download uri.",
+              e);
+    }
+  }
+
+  String makeSearchModelVersions(String searchFilter,
+                                 int maxResults,
+                                 List<String> orderBy,
+                                 String pageToken) {
+    try {
+      URIBuilder builder = new URIBuilder("model-versions/search")
+              .addParameter("max_results", Integer.toString(maxResults));
+      if (searchFilter != null && searchFilter != "") {
+        builder.addParameter("filter", searchFilter);
+      }
+      if (pageToken != null && pageToken != "") {
+        builder.addParameter("page_token", pageToken);
+      }
+      for( String order: orderBy) {
+        builder.addParameter("order_by", order);
+      }
+      return builder.build().toString();
+    } catch (URISyntaxException e) {
+      throw new MlflowClientException(
+              "Failed to construct request URI for search model version.",
+              e);
+    }
+  }
+
+  String toJson(MessageOrBuilder mb) {
+    return print(mb);
+  }
+
+  GetExperiment.Response toGetExperimentResponse(String json) {
+    GetExperiment.Response.Builder builder = GetExperiment.Response.newBuilder();
+    merge(json, builder);
+    return builder.build();
+  }
+
+  GetExperimentByName.Response toGetExperimentByNameResponse(String json) {
+    GetExperimentByName.Response.Builder builder = GetExperimentByName.Response.newBuilder();
+    merge(json, builder);
+    return builder.build();
+  }
+
+  SearchExperiments.Response toSearchExperimentsResponse(String json) {
+    SearchExperiments.Response.Builder builder = SearchExperiments.Response.newBuilder();
+    merge(json, builder);
+    return builder.build();
+  }
+
+  CreateExperiment.Response toCreateExperimentResponse(String json) {
+    CreateExperiment.Response.Builder builder = CreateExperiment.Response.newBuilder();
+    merge(json, builder);
+    return builder.build();
+  }
+
+  GetRun.Response toGetRunResponse(String json) {
+    GetRun.Response.Builder builder = GetRun.Response.newBuilder();
+    merge(json, builder);
+    return builder.build();
+  }
+
+  GetMetricHistory.Response toGetMetricHistoryResponse(String json) {
+    GetMetricHistory.Response.Builder builder = GetMetricHistory.Response.newBuilder();
+    merge(json, builder);
+    return builder.build();
+  }
+
+  CreateRun.Response toCreateRunResponse(String json) {
+    CreateRun.Response.Builder builder = CreateRun.Response.newBuilder();
+    merge(json, builder);
+    return builder.build();
+  }
+
+  SearchRuns.Response toSearchRunsResponse(String json) {
+    SearchRuns.Response.Builder builder = SearchRuns.Response.newBuilder();
+    merge(json, builder);
+    return builder.build();
+  }
+
+  GetLatestVersions.Response toGetLatestVersionsResponse(String json) {
+    GetLatestVersions.Response.Builder builder = GetLatestVersions.Response.newBuilder();
+    merge(json, builder);
+    return builder.build();
+  }
+
+  GetModelVersion.Response toGetModelVersionResponse(String json) {
+    GetModelVersion.Response.Builder builder = GetModelVersion.Response.newBuilder();
+    merge(json, builder);
+    return builder.build();
+  }
+
+  GetRegisteredModel.Response toGetRegisteredModelResponse(String json) {
+    GetRegisteredModel.Response.Builder builder = GetRegisteredModel.Response.newBuilder();
+    merge(json, builder);
+    return builder.build();
+  }
+
+  String toGetModelVersionDownloadUriResponse(String json) {
+    GetModelVersionDownloadUri.Response.Builder builder = GetModelVersionDownloadUri.Response
+            .newBuilder();
+    merge(json, builder);
+    return builder.getArtifactUri();
+  }
+
+  SearchModelVersions.Response toSearchModelVersionsResponse(String json) {
+    SearchModelVersions.Response.Builder builder = SearchModelVersions.Response.newBuilder();
+    merge(json, builder);
+    return builder.build();
+  }
+
+  private String print(MessageOrBuilder message) {
+    try {
+      return JsonFormat.printer().preservingProtoFieldNames().print(message);
+    } catch (InvalidProtocolBufferException e) {
+      throw new MlflowClientException("Failed to serialize message " + message, e);
+    }
+  }
+
+  private void merge(String json, com.google.protobuf.Message.Builder builder) {
+    try {
+      JsonFormat.parser().ignoringUnknownFields().merge(json, builder);
+    } catch (InvalidProtocolBufferException e) {
+      throw new MlflowClientException("Failed to serialize json " + json + " into " + builder, e);
+    }
+  }
+}

mlflow/java/client/src/main/java/org/mlflow/tracking/ModelVersionsPage.java

@@ -0,0 +1,83 @@
+package org.mlflow.tracking;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.Optional;
+import org.mlflow.api.proto.ModelRegistry.*;
+
+public class ModelVersionsPage implements Page<ModelVersion> {
+
+  private final String token;
+  private final List<ModelVersion> mvs;
+
+  private final MlflowClient client;
+  private final String searchFilter;
+  private final List<String> orderBy;
+  private final int maxResults;
+
+  /**
+   * Creates a fixed size page of ModelVersions.
+   */
+  ModelVersionsPage(List<ModelVersion> mvs,
+                    String token,
+                    String searchFilter,
+                    int maxResults,
+                    List<String> orderBy,
+                    MlflowClient client) {
+    this.mvs = Collections.unmodifiableList(mvs);
+    this.token = token;
+    this.searchFilter = searchFilter;
+    this.orderBy = orderBy;
+    this.maxResults = maxResults;
+    this.client = client;
+  }
+
+  /**
+   * @return The number of model versions in the page.
+   */
+  public int getPageSize() {
+    return this.mvs.size();
+  }
+
+  /**
+   * @return True if a token for the next page exists and isn't empty. Otherwise returns false.
+   */
+  public boolean hasNextPage() {
+    return this.token != null && this.token != "";
+  }
+
+  /**
+   * @return An optional with the token for the next page.
+   * Empty if the token doesn't exist or is empty.
+   */
+  public Optional<String> getNextPageToken() {
+    if (this.hasNextPage()) {
+      return Optional.of(this.token);
+    } else {
+      return Optional.empty();
+    }
+  }
+
+  /**
+   * @return The next page of model versions matching the search criteria.
+   * If there are no more pages, an {@link org.mlflow.tracking.EmptyPage} will be returned.
+   */
+  public Page<ModelVersion> getNextPage() {
+    if (this.hasNextPage()) {
+      return this.client.searchModelVersions(this.searchFilter,
+                                             this.maxResults,
+                                             this.orderBy,
+                                             this.token);
+    } else {
+      return new EmptyPage();
+    }
+  }
+
+  /**
+   * @return An iterable over the model versions in this page.
+   */
+  public List<ModelVersion> getItems() {
+    return mvs;
+  }
+
+}

mlflow/java/client/src/main/java/org/mlflow/tracking/Page.java

@@ -0,0 +1,35 @@
+package org.mlflow.tracking;
+
+import java.lang.Iterable;
+import java.util.Optional;
+
+public interface Page<E> {
+
+  /**
+   * @return The number of elements in this page.
+   */
+  public int getPageSize();
+
+  /**
+   * @return True if there are more pages that can be retrieved from the API.
+   */
+  public boolean hasNextPage();
+
+  /**
+   * @return An Optional of the token string to get the next page. 
+   * Empty if there is no next page.
+   */
+  public Optional<String> getNextPageToken();
+
+  /**
+   * @return Retrieves the next Page object using the next page token,
+   * or returns an empty page if there are no more pages.
+   */
+  public Page<E> getNextPage();
+
+  /**
+   * @return A List of the elements in this Page.
+   */
+  public Iterable<E> getItems();
+
+}

mlflow/java/client/src/main/java/org/mlflow/tracking/RunsPage.java

@@ -0,0 +1,92 @@
+package org.mlflow.tracking;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.ArrayList;
+import java.util.Optional;
+import org.mlflow.api.proto.Service.*;
+
+public class RunsPage implements Page<Run> {
+
+  private final String token;
+  private final List<Run> runs;
+
+  private final MlflowClient client;
+  private final List<String> experimentIds;
+  private final String searchFilter;
+  private final ViewType runViewType;
+  private final List<String> orderBy;
+  private final int maxResults;
+
+  /**
+   * Creates a fixed size page of Runs.
+   */
+  RunsPage(List<Run> runs,
+                  String token,
+                  List<String> experimentIds,
+                  String searchFilter,
+                  ViewType runViewType,
+                  int maxResults,
+                  List<String> orderBy,
+                  MlflowClient client) {
+    this.runs = Collections.unmodifiableList(runs);
+    this.token = token;
+    this.experimentIds = experimentIds;
+    this.searchFilter = searchFilter;
+    this.runViewType = runViewType;
+    this.orderBy = orderBy;
+    this.maxResults = maxResults;
+    this.client = client;
+  }
+
+  /**
+   * @return The number of runs in the page.
+   */
+  public int getPageSize() {
+    return this.runs.size();
+  }
+
+  /**
+   * @return True if a token for the next page exists and isn't empty. Otherwise returns false.
+   */
+  public boolean hasNextPage() {
+    return this.token != null && this.token != "";
+  }
+
+  /**
+   * @return An optional with the token for the next page. 
+   * Empty if the token doesn't exist or is empty.
+   */
+  public Optional<String> getNextPageToken() {
+    if (this.hasNextPage()) {
+      return Optional.of(this.token);
+    } else {
+      return Optional.empty();
+    }
+  }
+
+  /**
+   * @return The next page of runs matching the search criteria. 
+   * If there are no more pages, an {@link org.mlflow.tracking.EmptyPage} will be returned.
+   */
+  public Page<Run> getNextPage() {
+    if (this.hasNextPage()) {
+      return this.client.searchRuns(this.experimentIds,
+                                    this.searchFilter,
+                                    this.runViewType,
+                                    this.maxResults,
+                                    this.orderBy,
+                                    this.token);
+    } else {
+      return new EmptyPage();
+    }
+  }
+
+  /**
+   * @return An iterable over the runs in this page.
+   */
+  public List<Run> getItems() {
+    return runs;
+  }
+
+}

mlflow/java/client/src/main/java/org/mlflow/tracking/creds/BasicMlflowHostCreds.java

@@ -0,0 +1,73 @@
+package org.mlflow.tracking.creds;
+
+/** A static hostname and optional credentials to talk to an MLflow server. */
+public class BasicMlflowHostCreds implements MlflowHostCreds, MlflowHostCredsProvider {
+  private String host;
+  private String username;
+  private String password;
+  private String token;
+  private boolean shouldIgnoreTlsVerification;
+
+  public BasicMlflowHostCreds(String host) {
+    this.host = host;
+  }
+
+  public BasicMlflowHostCreds(String host, String username, String password) {
+    this.host = host;
+    this.username = username;
+    this.password = password;
+  }
+
+  public BasicMlflowHostCreds(String host, String token) {
+    this.host = host;
+    this.token = token;
+  }
+
+  public BasicMlflowHostCreds(
+      String host,
+      String username,
+      String password,
+      String token,
+      boolean shouldIgnoreTlsVerification) {
+    this.host = host;
+    this.username = username;
+    this.password = password;
+    this.token = token;
+    this.shouldIgnoreTlsVerification = shouldIgnoreTlsVerification;
+  }
+
+  @Override
+  public String getHost() {
+    return host;
+  }
+
+  @Override
+  public String getUsername() {
+    return username;
+  }
+
+  @Override
+  public String getPassword() {
+    return password;
+  }
+
+  @Override
+  public String getToken() {
+    return token;
+  }
+
+  @Override
+  public boolean shouldIgnoreTlsVerification() {
+    return shouldIgnoreTlsVerification;
+  }
+
+  @Override
+  public MlflowHostCreds getHostCreds() {
+    return this;
+  }
+
+  @Override
+  public void refresh() {
+    // no-op
+  }
+}

mlflow/java/client/src/main/java/org/mlflow/tracking/creds/DatabricksConfigHostCredsProvider.java

@@ -0,0 +1,103 @@
+package org.mlflow.tracking.creds;
+
+import java.io.File;
+import java.io.FileReader;
+import java.io.IOException;
+import java.nio.file.Paths;
+
+import org.apache.commons.configuration2.INIConfiguration;
+import org.apache.commons.configuration2.SubnodeConfiguration;
+import org.apache.commons.configuration2.ex.ConfigurationException;
+
+public class DatabricksConfigHostCredsProvider extends DatabricksHostCredsProvider {
+  private static final String CONFIG_FILE_ENV_VAR = "DATABRICKS_CONFIG_FILE";
+
+  private final String profile;
+
+  private DatabricksMlflowHostCreds hostCreds;
+
+  public DatabricksConfigHostCredsProvider(String profile) {
+    this.profile = profile;
+  }
+
+  public DatabricksConfigHostCredsProvider() {
+    this.profile = null;
+  }
+
+  private void loadConfigIfNecessary() {
+    if (hostCreds == null) {
+      reloadConfig();
+    }
+  }
+
+  private void reloadConfig() {
+    String basePath = System.getenv(CONFIG_FILE_ENV_VAR);
+    if (basePath == null) {
+      String userHome = System.getProperty("user.home");
+      basePath = Paths.get(userHome, ".databrickscfg").toString();
+    }
+
+    if (!new File(basePath).isFile()) {
+      throw new IllegalStateException("Could not find Databricks configuration file" +
+        " (" + basePath + "). Please run 'databricks configure' using the Databricks CLI.");
+    }
+
+    INIConfiguration ini;
+    try {
+      ini = new INIConfiguration();
+      try (FileReader reader = new FileReader(basePath)) {
+        ini.read(reader);
+      }
+    } catch (IOException | ConfigurationException e) {
+      throw new IllegalStateException("Failed to load databrickscfg file at " + basePath, e);
+    }
+
+    SubnodeConfiguration section;
+    if (profile == null) {
+      section = ini.getSection("DEFAULT");
+      if (section == null || section.isEmpty()) {
+        throw new IllegalStateException("Could not find 'DEFAULT' section within config file" +
+          " (" + basePath + "). Please run 'databricks configure' using the Databricks CLI.");
+      }
+    } else {
+      section = ini.getSection(profile);
+      if (section == null || section.isEmpty()) {
+        throw new IllegalStateException("Could not find '" + profile + "' section within config" +
+          " file  (" + basePath + "). Please run 'databricks configure --profile " + profile + "'" +
+          " using the Databricks CLI.");
+      }
+    }
+    assert (section != null);
+
+    String host = section.getString("host");
+    String username = section.getString("username");
+    String password = section.getString("password");
+    String token = section.getString("token");
+    boolean insecure = "true".equalsIgnoreCase(section.getString("insecure", "false"));
+
+    if (host == null) {
+      throw new IllegalStateException("No 'host' configured within Databricks config file" +
+        " (" + basePath + "). Please run 'databricks configure' using the Databricks CLI.");
+    }
+
+    boolean hasValidUserPassword = username != null && password != null;
+    boolean hasValidToken = token != null;
+    if (!hasValidUserPassword && !hasValidToken) {
+      throw new IllegalStateException("No authentication configured within Databricks config file" +
+        " (" + basePath + "). Please run 'databricks configure' using the Databricks CLI.");
+    }
+
+    this.hostCreds = new DatabricksMlflowHostCreds(host, username, password, token, insecure);
+  }
+
+  @Override
+  public DatabricksMlflowHostCreds getHostCreds() {
+    loadConfigIfNecessary();
+    return hostCreds;
+  }
+
+  @Override
+  public void refresh() {
+    reloadConfig();
+  }
+}

mlflow/java/client/src/main/java/org/mlflow/tracking/creds/DatabricksDynamicHostCredsProvider.java

@@ -0,0 +1,49 @@
+package org.mlflow.tracking.creds;
+
+import java.util.Map;
+
+import com.google.common.annotations.VisibleForTesting;
+import org.mlflow.tracking.utils.DatabricksContext;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class DatabricksDynamicHostCredsProvider extends DatabricksHostCredsProvider {
+  private static final Logger logger = LoggerFactory.getLogger(
+    DatabricksDynamicHostCredsProvider.class);
+
+  private final Map<String, String> configProvider;
+
+  private DatabricksDynamicHostCredsProvider(Map<String, String> configProvider) {
+    this.configProvider = configProvider;
+  }
+
+  public static DatabricksDynamicHostCredsProvider createIfAvailable() {
+    return createIfAvailable(DatabricksContext.CONFIG_PROVIDER_CLASS_NAME);
+  }
+
+  @VisibleForTesting
+  static DatabricksDynamicHostCredsProvider createIfAvailable(String className) {
+    Map<String, String> configProvider =
+      DatabricksContext.getConfigProviderIfAvailable(className);
+    if (configProvider == null) {
+      return null;
+    }
+    return new DatabricksDynamicHostCredsProvider(configProvider);
+  }
+
+  @Override
+  public DatabricksMlflowHostCreds getHostCreds() {
+    return new DatabricksMlflowHostCreds(
+      configProvider.get("host"),
+      configProvider.get("username"),
+      configProvider.get("password"),
+      configProvider.get("token"),
+      "true".equals(configProvider.get("shouldIgnoreTlsVerification"))
+    );
+  }
+
+  @Override
+  public void refresh() {
+    // no-op
+  }
+}

mlflow/java/client/src/main/java/org/mlflow/tracking/creds/DatabricksHostCredsProvider.java

@@ -0,0 +1,11 @@
+package org.mlflow.tracking.creds;
+
+abstract class DatabricksHostCredsProvider implements MlflowHostCredsProvider {
+
+  @Override
+  public abstract DatabricksMlflowHostCreds getHostCreds();
+
+  @Override
+  public abstract void refresh();
+
+}

mlflow/java/client/src/main/java/org/mlflow/tracking/creds/DatabricksMlflowHostCreds.java

@@ -0,0 +1,22 @@
+package org.mlflow.tracking.creds;
+
+/** Credentials to talk to a Databricks-hosted MLflow server. */
+public final class DatabricksMlflowHostCreds extends BasicMlflowHostCreds {
+
+  public DatabricksMlflowHostCreds(String host, String username, String password) {
+    super(host, username, password);
+  }
+
+  public DatabricksMlflowHostCreds(String host, String token) {
+    super(host, token);
+  }
+
+  public DatabricksMlflowHostCreds(
+      String host,
+      String username,
+      String password,
+      String token,
+      boolean shouldIgnoreTlsVerification) {
+    super(host, username, password, token, shouldIgnoreTlsVerification);
+  }
+}

mlflow/java/client/src/main/java/org/mlflow/tracking/creds/HostCredsProviderChain.java

@@ -0,0 +1,48 @@
+package org.mlflow.tracking.creds;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.mlflow.tracking.MlflowClientException;
+
+public class HostCredsProviderChain implements MlflowHostCredsProvider {
+  private static final Logger logger = LoggerFactory.getLogger(HostCredsProviderChain.class);
+
+  private final List<MlflowHostCredsProvider> hostCredsProviders = new ArrayList<>();
+
+  public HostCredsProviderChain(MlflowHostCredsProvider... hostCredsProviders) {
+    this.hostCredsProviders.addAll(Arrays.asList(hostCredsProviders));
+  }
+
+  @Override
+  public MlflowHostCreds getHostCreds() {
+    List<String> exceptionMessages = new ArrayList<>();
+    for (MlflowHostCredsProvider provider : hostCredsProviders) {
+      try {
+        MlflowHostCreds hostCreds = provider.getHostCreds();
+
+        if (hostCreds != null && hostCreds.getHost() != null) {
+          logger.debug("Loading credentials from " + provider.toString());
+          return hostCreds;
+        }
+      } catch (Exception e) {
+        String message = provider + ": " + e.getMessage();
+        logger.debug("Unable to load credentials from " + message);
+        exceptionMessages.add(message);
+      }
+    }
+    throw new MlflowClientException("Unable to load MLflow Host/Credentials from any provider in" +
+      " the chain: " + exceptionMessages);
+  }
+
+  @Override
+  public void refresh() {
+    for (MlflowHostCredsProvider provider : hostCredsProviders) {
+      provider.refresh();
+    }
+  }
+}

mlflow/java/client/src/main/java/org/mlflow/tracking/creds/MlflowHostCreds.java

@@ -0,0 +1,33 @@
+package org.mlflow.tracking.creds;
+
+/**
+ * Provides a hostname and optional authentication for talking to an MLflow server.
+ */
+public interface MlflowHostCreds {
+  /** Hostname (e.g., http://localhost:5000) to MLflow server. */
+  String getHost();
+
+  /**
+   * Username to use with Basic authentication when talking to server.
+   * If this is specified, password must also be specified.
+   */
+  String getUsername();
+
+  /**
+   * Password to use with Basic authentication when talking to server.
+   * If this is specified, username must also be specified.
+   */
+  String getPassword();
+
+  /**
+   * Token to use with Bearer authentication when talking to server.
+   * If provided, user/password authentication will be ignored.
+   */
+  String getToken();
+
+  /**
+   * If true, we will not verify the server's hostname or TLS certificate.
+   * This is useful for certain testing situations, but should never be true in production.
+   */
+  boolean shouldIgnoreTlsVerification();
+}

mlflow/java/client/src/main/java/org/mlflow/tracking/creds/MlflowHostCredsProvider.java

@@ -0,0 +1,11 @@
+package org.mlflow.tracking.creds;
+
+/** Provides a dynamic, refreshable set of MlflowHostCreds. */
+public interface MlflowHostCredsProvider {
+
+  /** Returns a valid MlflowHostCreds. This may be cached. */
+  MlflowHostCreds getHostCreds();
+
+  /** Refreshes the underlying credentials. May be a no-op. */
+  void refresh();
+}

mlflow/java/client/src/main/java/org/mlflow/tracking/creds/package-info.java

@@ -0,0 +1,2 @@
+/** Support for custom tracking service discovery and authentication. */
+package org.mlflow.tracking.creds;

mlflow/java/client/src/main/java/org/mlflow/tracking/package-info.java

@@ -0,0 +1,5 @@
+/**
+ * MLflow Tracking provides a Java CRUD interface to MLflow Experiments and Runs --
+ * to create and log to MLflow runs, use the {@link org.mlflow.tracking.MlflowContext} interface.
+ */
+package org.mlflow.tracking;

mlflow/java/client/src/main/java/org/mlflow/tracking/samples/FluentExample.java

@@ -0,0 +1,61 @@
+package org.mlflow.tracking.samples;
+
+import com.google.common.collect.ImmutableMap;
+import org.mlflow.tracking.ActiveRun;
+import org.mlflow.tracking.MlflowContext;
+
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+
+public class FluentExample {
+  public static void main(String[] args) {
+    MlflowContext mlflow = new MlflowContext();
+    ExecutorService executor = Executors.newFixedThreadPool(10);
+
+    // Vanilla usage
+    {
+      ActiveRun run = mlflow.startRun("run");
+      run.logParam("alpha", "0.0");
+      run.logMetric("MSE", 0.0);
+      run.setTags(ImmutableMap.of(
+        "company", "databricks",
+        "org", "engineering"
+      ));
+      run.endRun();
+    }
+
+    // Lambda usage
+    {
+      mlflow.withActiveRun("lambda run", (activeRun -> {
+        activeRun.logParam("layers", "4");
+        // Perform training code
+      }));
+    }
+    // Log one parent run and 5 children run
+    {
+      ActiveRun run = mlflow.startRun("parent run");
+      for (int i = 0; i <= 5; i++) {
+        ActiveRun childRun = mlflow.startRun("child run", run.getId());
+        childRun.logParam("iteration", Integer.toString(i));
+        childRun.endRun();
+      }
+      run.endRun();
+    }
+
+    // Log one parent run and 5 children run (multithreaded)
+    {
+      ActiveRun run = mlflow.startRun("parent run (multithreaded)");
+      for (int i = 0; i <= 5; i++) {
+        final int i0 = i;
+        executor.submit(() -> {
+          ActiveRun childRun = mlflow.startRun("child run (multithreaded)", run.getId());
+          childRun.logParam("iteration", Integer.toString(i0));
+          childRun.endRun();
+        });
+      }
+      run.endRun();
+    }
+    executor.shutdown();
+    mlflow.getClient().close();
+  }
+}

mlflow/java/client/src/main/java/org/mlflow/tracking/samples/QuickStartDriver.java

@@ -0,0 +1,78 @@
+package org.mlflow.tracking.samples;
+
+import java.util.List;
+import java.util.Optional;
+
+import org.mlflow.api.proto.Service.*;
+import org.mlflow.tracking.MlflowClient;
+
+/**
+ * This is an example application which uses the MLflow Tracking API to create and manage
+ * experiments and runs.
+ */
+public class QuickStartDriver {
+  public static void main(String[] args) throws Exception {
+    (new QuickStartDriver()).process(args);
+  }
+
+  void process(String[] args) throws Exception {
+    MlflowClient client;
+    if (args.length < 1) {
+      client = new MlflowClient();
+    } else {
+      client = new MlflowClient(args[0]);
+    }
+
+    System.out.println("====== createExperiment");
+    String expName = "Exp_" + System.currentTimeMillis();
+    String expId = client.createExperiment(expName);
+    System.out.println("createExperiment: expId=" + expId);
+
+    System.out.println("====== getExperiment");
+    Experiment exp = client.getExperiment(expId);
+    System.out.println("getExperiment: " + exp);
+
+    System.out.println("====== searchExperiments");
+    List<Experiment> exps = client.searchExperiments().getItems();
+    System.out.println("#experiments: " + exps.size());
+    exps.forEach(e -> System.out.println("  Exp: " + e));
+
+    createRun(client, expId);
+
+    System.out.println("====== getExperiment again");
+    Experiment exp2 = client.getExperiment(expId);
+    System.out.println("getExperiment: " + exp2);
+
+    System.out.println("====== getExperiment by name");
+    Optional<Experiment> exp3 = client.getExperimentByName(expName);
+    System.out.println("getExperimentByName: " + exp3);
+    client.close();
+  }
+
+  void createRun(MlflowClient client, String expId) {
+    System.out.println("====== createRun");
+
+    // Create run
+    String sourceFile = "MyFile.java";
+
+    RunInfo runCreated = client.createRun(expId);
+    System.out.println("CreateRun: " + runCreated);
+    String runId = runCreated.getRunUuid();
+
+    // Log parameters
+    client.logParam(runId, "min_samples_leaf", "2");
+    client.logParam(runId, "max_depth", "3");
+
+    // Log metrics
+    client.logMetric(runId, "auc", 2.12F);
+    client.logMetric(runId, "accuracy_score", 3.12F);
+    client.logMetric(runId, "zero_one_loss", 4.12F);
+
+    // Update finished run
+    client.setTerminated(runId, RunStatus.FINISHED);
+
+    // Get run details
+    Run run = client.getRun(runId);
+    System.out.println("GetRun: " + run);
+  }
+}

mlflow/java/client/src/main/java/org/mlflow/tracking/utils/DatabricksContext.java

@@ -0,0 +1,131 @@
+package org.mlflow.tracking.utils;
+
+import com.google.common.annotations.VisibleForTesting;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class DatabricksContext {
+  public static final String CONFIG_PROVIDER_CLASS_NAME =
+    "com.databricks.config.DatabricksClientSettingsProvider";
+  private static final Logger logger = LoggerFactory.getLogger(
+    DatabricksContext.class);
+  private final Map<String, String> configProvider;
+
+  private DatabricksContext(Map<String, String> configProvider) {
+    this.configProvider = configProvider;
+  }
+
+  public static DatabricksContext createIfAvailable() {
+    return createIfAvailable(CONFIG_PROVIDER_CLASS_NAME);
+  }
+
+  @VisibleForTesting
+  static DatabricksContext createIfAvailable(String className) {
+    Map<String, String> configProvider = getConfigProviderIfAvailable(className);
+    if (configProvider == null) {
+      return null;
+    }
+    return new DatabricksContext(configProvider);
+  }
+
+  public Map<String, String> getTags() {
+    if (isInDatabricksNotebook()) {
+      return getTagsForDatabricksNotebook();
+    } else if (isInDatabricksJob()) {
+      return getTagsForDatabricksJob();
+    } else {
+      return new HashMap<>();
+    }
+  }
+
+  public boolean isInDatabricksNotebook() {
+    return configProvider.get("notebookId") != null;
+  }
+
+  /**
+   * Should only be called if isInDatabricksNotebook() is true.
+   */
+  private Map<String, String> getTagsForDatabricksNotebook() {
+    Map<String, String> tagsForNotebook = new HashMap<>();
+    String notebookId = getNotebookId();
+    if (notebookId != null) {
+      tagsForNotebook.put(MlflowTagConstants.DATABRICKS_NOTEBOOK_ID, notebookId);
+    }
+    String notebookPath = configProvider.get("notebookPath");
+    if (notebookPath != null) {
+      tagsForNotebook.put(MlflowTagConstants.SOURCE_NAME, notebookPath);
+      tagsForNotebook.put(MlflowTagConstants.DATABRICKS_NOTEBOOK_PATH, notebookPath);
+      tagsForNotebook.put(MlflowTagConstants.SOURCE_TYPE, "NOTEBOOK");
+    }
+    String webappUrl = configProvider.get("host");
+    if (webappUrl != null) {
+      tagsForNotebook.put(MlflowTagConstants.DATABRICKS_WEBAPP_URL, webappUrl);
+    }
+    return tagsForNotebook;
+  }
+
+  /**
+   * Should only be called if isInDatabricksNotebook() is true.
+   */
+  public String getNotebookId() {
+    if (!isInDatabricksNotebook()) {
+      throw new IllegalArgumentException(
+        "getNotebookId() should not be called when isInDatabricksNotebook() is false"
+      );
+    }
+    return configProvider.get("notebookId");
+  }
+
+  public String getNotebookPath() {
+    if (!isInDatabricksNotebook()) {
+      throw new IllegalArgumentException(
+        "getNotebookPath() should not be called when isInDatabricksNotebook() is false"
+      );
+    }
+    return configProvider.get("notebookPath");
+  }
+
+  private boolean isInDatabricksJob() {
+    return configProvider.get("jobId") != null;
+  }
+
+  /**
+   * Should only be called if isInDatabricksJob() is true.
+   */
+  private Map<String, String> getTagsForDatabricksJob() {
+    Map<String, String> tagsForJob = new HashMap<>();
+    String jobId = configProvider.get("jobId");
+    String jobRunId = configProvider.get("jobRunId");
+    String jobType = configProvider.get("jobType");
+    String webappUrl = configProvider.get("host");
+    if (jobId != null && jobRunId != null) {
+      tagsForJob.put(MlflowTagConstants.DATABRICKS_JOB_ID, jobId);
+      tagsForJob.put(MlflowTagConstants.DATABRICKS_JOB_RUN_ID, jobRunId);
+      tagsForJob.put(MlflowTagConstants.SOURCE_TYPE, "JOB");
+      tagsForJob.put(MlflowTagConstants.SOURCE_NAME,
+                          String.format("jobs/%s/run/%s", jobId, jobRunId));
+    }
+    if (jobType != null) {
+      tagsForJob.put(MlflowTagConstants.DATABRICKS_JOB_TYPE, jobType);
+    }
+    if (webappUrl != null) {
+      tagsForJob.put(MlflowTagConstants.DATABRICKS_WEBAPP_URL, webappUrl);
+    }
+    return tagsForJob;
+  }
+
+  public static Map<String, String> getConfigProviderIfAvailable(String className) {
+    try {
+      Class<?> cls = Class.forName(className);
+      return (Map<String, String>) cls.newInstance();
+    } catch (ClassNotFoundException e) {
+      return null;
+    } catch (IllegalAccessException | InstantiationException e) {
+      logger.warn("Found but failed to invoke dynamic config provider", e);
+      return null;
+    }
+  }
+}

mlflow/java/client/src/main/java/org/mlflow/tracking/utils/MlflowTagConstants.java

@@ -0,0 +1,19 @@
+package org.mlflow.tracking.utils;
+
+public class MlflowTagConstants {
+  public static final String PARENT_RUN_ID = "mlflow.parentRunId";
+  public static final String RUN_NAME = "mlflow.runName";
+  public static final String USER = "mlflow.user";
+  public static final String SOURCE_TYPE = "mlflow.source.type";
+  public static final String SOURCE_NAME = "mlflow.source.name";
+  public static final String DATABRICKS_NOTEBOOK_ID = "mlflow.databricks.notebookID";
+  public static final String DATABRICKS_NOTEBOOK_PATH = "mlflow.databricks.notebookPath";
+  // The JOB_ID, JOB_RUN_ID, and JOB_TYPE tags are used for automatically recording Job 
+  // information when MLflow Tracking APIs are used within a Databricks Job
+  public static final String DATABRICKS_JOB_ID = "mlflow.databricks.jobID";
+  public static final String DATABRICKS_JOB_RUN_ID = "mlflow.databricks.jobRunID";
+  public static final String DATABRICKS_JOB_TYPE = "mlflow.databricks.jobType";
+  public static final String DATABRICKS_WEBAPP_URL = "mlflow.databricks.webappURL";
+  public static final String MLFLOW_EXPERIMENT_SOURCE_TYPE = "mlflow.experiment.sourceType";
+  public static final String MLFLOW_EXPERIMENT_SOURCE_ID = "mlflow.experiment.sourceId";
+}

mlflow/java/client/src/main/resources/log4j.properties

@@ -0,0 +1,5 @@
+log4j.rootLogger=info, console
+log4j.logger.com.databricks=info
+log4j.appender.console=org.apache.log4j.ConsoleAppender
+log4j.appender.console.layout=org.apache.log4j.PatternLayout
+log4j.appender.console.layout.ConversionPattern=%p  %c.%M:%L: %m%n

mlflow/java/client/src/test/java/org/mlflow/artifacts/CliBasedArtifactRepositoryTest.java

@@ -0,0 +1,237 @@
+package org.mlflow.artifacts;
+
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+import com.google.common.collect.Sets;
+import org.apache.commons.io.FileUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.testng.Assert;
+import org.testng.annotations.AfterSuite;
+import org.testng.annotations.BeforeSuite;
+import org.testng.annotations.Test;
+
+import org.mlflow.api.proto.Service.FileInfo;
+import org.mlflow.api.proto.Service.RunInfo;
+import org.mlflow.tracking.MlflowClient;
+import org.mlflow.tracking.TestClientProvider;
+import org.mlflow.tracking.creds.BasicMlflowHostCreds;
+import org.mlflow.tracking.creds.DatabricksMlflowHostCreds;
+import org.mlflow.tracking.creds.MlflowHostCreds;
+
+public class CliBasedArtifactRepositoryTest {
+  private static final Logger logger = LoggerFactory.getLogger(
+    CliBasedArtifactRepositoryTest.class);
+
+  private final TestClientProvider testClientProvider = new TestClientProvider();
+
+  private MlflowClient client;
+
+  @BeforeSuite
+  public void beforeAll() throws IOException {
+    client = testClientProvider.initializeClientAndServer();
+  }
+
+  @AfterSuite
+  public void afterAll() throws InterruptedException {
+    testClientProvider.cleanupClientAndServer();
+  }
+
+  private CliBasedArtifactRepository newRepo() {
+    RunInfo runInfo = client.createRun();
+    logger.info("Created run with id=" + runInfo.getRunUuid() + " and artifactUri=" +
+      runInfo.getArtifactUri());
+    return new CliBasedArtifactRepository(runInfo.getArtifactUri(), runInfo.getRunUuid(),
+      testClientProvider.getClientHostCredsProvider(client));
+  }
+
+  @Test
+  public void testLogAndDownloadArtifact() throws IOException {
+    // Tests the logArtifact and downloadArtifacts APIs when targeting a single file.
+    ArtifactRepository repo = newRepo();
+    Path tempFile = Files.createTempFile(getClass().getSimpleName(), ".txt");
+    FileUtils.writeStringToFile(tempFile.toFile(), "Hello, World!", StandardCharsets.UTF_8);
+    repo.logArtifact(tempFile.toFile());
+    Path returnFile = repo.downloadArtifacts(tempFile.getFileName().toString()).toPath();
+    Assert.assertEquals(readFile(returnFile), "Hello, World!");
+  }
+
+  @Test
+  public void testLogListAndDownloadArtifacts() throws IOException {
+    // Tests the logArtifacts, list, downloadArtifacts APIs when targeting a set of files.
+    ArtifactRepository repo = newRepo();
+    Path tempDir = Files.createTempDirectory(getClass().getSimpleName());
+    FileUtils.writeStringToFile(tempDir.resolve("a").toFile(), "A", StandardCharsets.UTF_8);
+    FileUtils.writeStringToFile(tempDir.resolve("b").toFile(), "B", StandardCharsets.UTF_8);
+    repo.logArtifacts(tempDir.toFile());
+    Set<FileInfo> fileInfos = Sets.newHashSet(repo.listArtifacts());
+    Assert.assertEquals(fileInfos, Sets.newHashSet(fileInfo("a", 1), fileInfo("b", 1)));
+    Path returnDir = repo.downloadArtifacts().toPath();
+    Assert.assertEquals(readFile(returnDir.resolve("a")), "A");
+    Assert.assertEquals(readFile(returnDir.resolve("b")), "B");
+  }
+
+  @Test
+  public void testEverything() throws IOException {
+    // This is a comprehensive integration test which tests all 8 APIs exposed by ArtifactRepository
+    // on a mix of files, directories, and subdirectories.
+    ArtifactRepository repo = newRepo();
+
+    // Create a temporary directory with /childFile and /childDir/granchild as files.
+    Path tempDir = Files.createTempDirectory(getClass().getSimpleName());
+    Path childFile = tempDir.resolve("childFile");
+    String childContents = "File contents!";
+    FileUtils.writeStringToFile(childFile.toFile(), childContents, StandardCharsets.UTF_8);
+    Path childDir = tempDir.resolve("childDir");
+    Path grandchildFile = childDir.resolve("grandchild");
+    String grandchildContents = "Baby content!";
+    childDir.toFile().mkdir();
+    FileUtils.writeStringToFile(grandchildFile.toFile(), grandchildContents, StandardCharsets.UTF_8);
+
+    // Log artifacts such that we expect:
+    //   childFile
+    //   grandchild
+    //   subpath/childFile
+    //   subpath/subberpath/grandchild
+    repo.logArtifact(childFile.toFile());
+    repo.logArtifacts(childDir.toFile());
+    repo.logArtifact(childFile.toFile(), "subpath");
+    repo.logArtifacts(childDir.toFile(), "subpath/subberpath");
+
+    // List at the root, we should see childFile, grandchild, and subpath/.
+    List<FileInfo> fileInfos = repo.listArtifacts();
+    Set<FileInfo> expectedFileInfos = Sets.newHashSet(
+      fileInfo("childFile", childContents.length()),
+      fileInfo("grandchild", grandchildContents.length()),
+      dirInfo("subpath")
+    );
+    Assert.assertEquals(Sets.newHashSet(fileInfos), expectedFileInfos);
+
+    // List within subpath/, we should see childFile and subberpath.
+    List<FileInfo> subpathFileInfos = repo.listArtifacts("subpath");
+    Set<FileInfo> expectedSubpathFileInfos = Sets.newHashSet(
+      fileInfo("subpath/childFile", childContents.length()),
+      dirInfo("subpath/subberpath")
+    );
+    Assert.assertEquals(Sets.newHashSet(subpathFileInfos), expectedSubpathFileInfos);
+
+    // Download everything, and confirm that we have the four expected files.
+    Path allArtifacts = repo.downloadArtifacts().toPath();
+    Assert.assertEquals(childContents, readFile(allArtifacts.resolve("childFile")));
+    Assert.assertEquals(childContents, readFile(allArtifacts.resolve("subpath/childFile")));
+    Assert.assertEquals(grandchildContents, readFile(allArtifacts.resolve("grandchild")));
+    Assert.assertEquals(grandchildContents,
+      readFile(allArtifacts.resolve("subpath/subberpath/grandchild")));
+
+    // Download subpath/subberpath, and confirm that we have just the grandchild.
+    Path subberpathArtifacts = repo.downloadArtifacts("subpath/subberpath").toPath();
+    Assert.assertEquals(grandchildContents, readFile(subberpathArtifacts.resolve("grandchild")));
+    Assert.assertEquals(subberpathArtifacts.toFile().list(), new String[] {"grandchild"});
+  }
+
+  @Test
+  public void testSettingProcessEnvBasic() {
+    CliBasedArtifactRepository repo = newRepo();
+    MlflowHostCreds hostCreds = new BasicMlflowHostCreds("just-host");
+    Map<String, String> env = new HashMap<>();
+    repo.setProcessEnvironment(env, hostCreds);
+    Map<String, String> expectedEnv = new HashMap<>();
+    expectedEnv.put("MLFLOW_TRACKING_URI", "just-host");
+    Assert.assertEquals(env, expectedEnv);
+  }
+
+  @Test
+  public void testSettingProcessEnvUserPass() {
+    CliBasedArtifactRepository repo = newRepo();
+    MlflowHostCreds hostCreds = new BasicMlflowHostCreds("just-host", "user", "pass");
+    Map<String, String> env = new HashMap<>();
+    repo.setProcessEnvironment(env, hostCreds);
+    Map<String, String> expectedEnv = new HashMap<>();
+    expectedEnv.put("MLFLOW_TRACKING_URI", "just-host");
+    expectedEnv.put("MLFLOW_TRACKING_USERNAME", "user");
+    expectedEnv.put("MLFLOW_TRACKING_PASSWORD", "pass");
+    Assert.assertEquals(env, expectedEnv);
+  }
+
+  @Test
+  public void testSettingProcessEnvToken() {
+    CliBasedArtifactRepository repo = newRepo();
+    MlflowHostCreds hostCreds = new BasicMlflowHostCreds("just-host", "token");
+    Map<String, String> env = new HashMap<>();
+    repo.setProcessEnvironment(env, hostCreds);
+    Map<String, String> expectedEnv = new HashMap<>();
+    expectedEnv.put("MLFLOW_TRACKING_URI", "just-host");
+    expectedEnv.put("MLFLOW_TRACKING_TOKEN", "token");
+    Assert.assertEquals(env, expectedEnv);
+  }
+
+  @Test
+  public void testSettingProcessEnvInsecure() {
+    CliBasedArtifactRepository repo = newRepo();
+    MlflowHostCreds hostCreds = new BasicMlflowHostCreds("insecure-host", null, null, null,
+      true);
+    Map<String, String> env = new HashMap<>();
+    repo.setProcessEnvironment(env, hostCreds);
+    Map<String, String> expectedEnv = new HashMap<>();
+    expectedEnv.put("MLFLOW_TRACKING_URI", "insecure-host");
+    expectedEnv.put("MLFLOW_TRACKING_INSECURE_TLS", "true");
+    Assert.assertEquals(env, expectedEnv);
+  }
+
+  @Test
+  public void testSettingProcessEnvDatabricksUserPass() {
+    CliBasedArtifactRepository repo = newRepo();
+    DatabricksMlflowHostCreds hostCreds = new DatabricksMlflowHostCreds(
+      "just-host", "user", "pass");
+    Map<String, String> env = new HashMap<>();
+    repo.setProcessEnvironmentDatabricks(env, hostCreds);
+    Map<String, String> expectedEnv = new HashMap<>();
+    expectedEnv.put("DATABRICKS_HOST", "just-host");
+    expectedEnv.put("DATABRICKS_USERNAME", "user");
+    expectedEnv.put("DATABRICKS_PASSWORD", "pass");
+    Assert.assertEquals(env, expectedEnv);
+  }
+
+  @Test
+  public void testSettingProcessEnvDatabricksToken() {
+    CliBasedArtifactRepository repo = newRepo();
+    DatabricksMlflowHostCreds hostCreds = new DatabricksMlflowHostCreds("just-host", "token");
+    Map<String, String> env = new HashMap<>();
+    repo.setProcessEnvironmentDatabricks(env, hostCreds);
+    Map<String, String> expectedEnv = new HashMap<>();
+    expectedEnv.put("DATABRICKS_HOST", "just-host");
+    expectedEnv.put("DATABRICKS_TOKEN", "token");
+    Assert.assertEquals(env, expectedEnv);
+  }
+
+  @Test
+  public void testSettingProcessEnvDatabricksInsecure() {
+    CliBasedArtifactRepository repo = newRepo();
+    DatabricksMlflowHostCreds hostCreds = new DatabricksMlflowHostCreds(
+      "insecure-host", null, null, null, true);
+    Map<String, String> env = new HashMap<>();
+    repo.setProcessEnvironmentDatabricks(env, hostCreds);
+    Map<String, String> expectedEnv = new HashMap<>();
+    expectedEnv.put("DATABRICKS_HOST", "insecure-host");
+    expectedEnv.put("DATABRICKS_INSECURE", "true");
+    Assert.assertEquals(env, expectedEnv);
+  }
+
+  private String readFile(Path path) throws IOException {
+    return FileUtils.readFileToString(path.toFile(), StandardCharsets.UTF_8);
+  }
+
+  private FileInfo fileInfo(String path, int fileSize) {
+    return FileInfo.newBuilder().setPath(path).setFileSize(fileSize).setIsDir(false).build();
+  }
+  private FileInfo dirInfo(String path) {
+    return FileInfo.newBuilder().setPath(path).setIsDir(true).build();
+  }
+}

mlflow/java/client/src/test/java/org/mlflow/tracking/ActiveRunTest.java

@@ -0,0 +1,174 @@
+package org.mlflow.tracking;
+
+import com.google.common.collect.ImmutableMap;
+import org.mlflow.api.proto.Service.*;
+import org.mockito.ArgumentCaptor;
+import org.testng.annotations.Test;
+import org.testng.collections.Lists;
+
+import java.io.File;
+import java.nio.file.Paths;
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+
+import static org.testng.Assert.*;
+import static org.mockito.Mockito.*;
+
+public class ActiveRunTest {
+
+  private static final String RUN_ID = "test-run-id";
+  private static final String ARTIFACT_URI = "dbfs:/artifact-uri";
+
+  private MlflowClient mockClient;
+
+  private ActiveRun getActiveRun() {
+    RunInfo r = RunInfo.newBuilder().setRunId(RUN_ID).setArtifactUri(ARTIFACT_URI).build();
+    this.mockClient = mock(MlflowClient.class);
+    return new ActiveRun(r, mockClient);
+  }
+
+  @Test
+  public void testGetId() {
+    assertEquals(getActiveRun().getId(), RUN_ID);
+  }
+
+  @Test
+  public void testLogParam() {
+    getActiveRun().logParam("param-key", "param-value");
+    verify(mockClient).logParam(RUN_ID, "param-key", "param-value");
+  }
+
+  @Test
+  public void testSetTag() {
+    getActiveRun().setTag("tag-key", "tag-value");
+    verify(mockClient).setTag(RUN_ID, "tag-key", "tag-value");
+  }
+
+  @Test
+  public void testLogMetric() {
+    getActiveRun().logMetric("metric-key", 1.0);
+    // The any is for the timestamp.
+    verify(mockClient).logMetric(eq(RUN_ID), eq("metric-key"), eq(1.0), anyLong(), eq(0L));
+  }
+
+  @Test
+  public void testLogMetricWithStep() {
+    getActiveRun().logMetric("metric-key", 1.0, 99);
+    // The any is for the timestamp.
+    verify(mockClient).logMetric(eq(RUN_ID), eq("metric-key"), eq(1.0), anyLong(), eq(99L));
+  }
+
+  @Test
+  public void testLogMetrics() {
+    ActiveRun activeRun = getActiveRun();
+    ArgumentCaptor<Iterable<Metric>> metricsArg = ArgumentCaptor.forClass(Iterable.class);
+    activeRun.logMetrics(ImmutableMap.of("a", 0.0, "b", 1.0));
+    verify(mockClient).logBatch(eq(RUN_ID), metricsArg.capture(), any(), any());
+
+    Set<Metric> metrics = new HashSet<>();
+    metricsArg.getValue().forEach(metrics::add);
+
+    assertTrue(metrics.stream()
+      .anyMatch(m -> m.getKey().equals("a") && m.getValue() == 0.0 && m.getStep() == 0));
+    assertTrue(metrics.stream()
+      .anyMatch(m -> m.getKey().equals("b") && m.getValue() == 1.0 && m.getStep() == 0));
+  }
+
+  @Test
+  public void testLogMetricsWithStep() {
+    ActiveRun activeRun = getActiveRun();
+    ArgumentCaptor<Iterable<Metric>> metricsArg = ArgumentCaptor.forClass(Iterable.class);
+    activeRun.logMetrics(ImmutableMap.of("a", 0.0, "b", 1.0), 99);
+    verify(mockClient).logBatch(eq(RUN_ID), metricsArg.capture(), any(), any());
+
+    Set<Metric> metrics = new HashSet<>();
+    metricsArg.getValue().forEach(metrics::add);
+
+    assertTrue(metrics.stream()
+      .anyMatch(m -> m.getKey().equals("a") && m.getValue() == 0.0 && m.getStep() == 99));
+    assertTrue(metrics.stream()
+      .anyMatch(m -> m.getKey().equals("b") && m.getValue() == 1.0 && m.getStep() == 99));
+  }
+
+  @Test
+  public void testLogParams() {
+    ActiveRun activeRun = getActiveRun();
+    ArgumentCaptor<Iterable<Param>> paramsArg = ArgumentCaptor.forClass(Iterable.class);
+    activeRun.logParams(ImmutableMap.of("a", "a", "b", "b"));
+    verify(mockClient).logBatch(eq(RUN_ID), any(), paramsArg.capture(), any());
+
+    Set<Param> params = new HashSet<>();
+    paramsArg.getValue().forEach(params::add);
+
+    assertTrue(params.stream()
+      .anyMatch(p -> p.getKey().equals("a") && p.getValue().equals("a")));
+    assertTrue(params.stream()
+      .anyMatch(p -> p.getKey().equals("b") && p.getValue().equals("b")));
+  }
+
+  @Test
+  public void testSetTags() {
+    ActiveRun activeRun = getActiveRun();
+    ArgumentCaptor<Iterable<RunTag>> tagsArg = ArgumentCaptor.forClass(Iterable.class);
+    activeRun.setTags(ImmutableMap.of("a", "a", "b", "b"));
+    verify(mockClient).logBatch(eq(RUN_ID), any(), any(), tagsArg.capture());
+
+    Set<RunTag> tags = new HashSet<>();
+    tagsArg.getValue().forEach(tags::add);
+
+    assertTrue(tags.stream()
+      .anyMatch(t -> t.getKey().equals("a") && t.getValue().equals("a")));
+    assertTrue(tags.stream()
+      .anyMatch(t -> t.getKey().equals("b") && t.getValue().equals("b")));
+  }
+
+  @Test
+  public void testLogArtifact() {
+    ActiveRun activeRun = getActiveRun();
+    activeRun.logArtifact(Paths.get("test"));
+    verify(mockClient).logArtifact(RUN_ID, new File("test"));
+  }
+
+  @Test
+  public void testLogArtifactWithArtifactPath() {
+    ActiveRun activeRun = getActiveRun();
+    activeRun.logArtifact(Paths.get("test"), "artifact-path");
+    verify(mockClient).logArtifact(RUN_ID, new File("test"), "artifact-path");
+  }
+
+  @Test
+  public void testLogArtifacts() {
+    ActiveRun activeRun = getActiveRun();
+    activeRun.logArtifacts(Paths.get("test"));
+    verify(mockClient).logArtifacts(RUN_ID, new File("test"));
+  }
+
+  @Test
+  public void testLogArtifactsWithArtifactPath() {
+    ActiveRun activeRun = getActiveRun();
+    activeRun.logArtifacts(Paths.get("test"), "artifact-path");
+    verify(mockClient).logArtifacts(RUN_ID, new File("test"), "artifact-path");
+  }
+
+  @Test
+  public void testGetArtifactUri() {
+    ActiveRun activeRun = getActiveRun();
+    assertEquals(activeRun.getArtifactUri(), ARTIFACT_URI);
+  }
+
+  @Test
+  public void testEndRun() {
+    ActiveRun activeRun = getActiveRun();
+    activeRun.endRun();
+    verify(mockClient).setTerminated(RUN_ID, RunStatus.FINISHED);
+  }
+
+  @Test
+  public void testEndRunWithStatus() {
+    ActiveRun activeRun = getActiveRun();
+    activeRun.endRun(RunStatus.FAILED);
+    verify(mockClient).setTerminated(RUN_ID, RunStatus.FAILED);
+  }
+}

mlflow/java/client/src/test/java/org/mlflow/tracking/HttpCallerTest.java

@@ -0,0 +1,153 @@
+package org.mlflow.tracking;
+
+
+import org.apache.http.HttpEntity;
+import org.apache.http.StatusLine;
+import org.apache.http.client.methods.CloseableHttpResponse;
+import org.apache.http.client.methods.HttpUriRequest;
+import org.apache.http.impl.client.CloseableHttpClient;
+import org.mlflow.tracking.creds.BasicMlflowHostCreds;
+import org.mlflow.tracking.creds.MlflowHostCreds;
+import org.mlflow.tracking.creds.MlflowHostCredsProvider;
+import org.testng.Assert;
+import org.testng.annotations.*;
+
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+
+import static org.mockito.Mockito.*;
+
+public class HttpCallerTest {
+
+  static CloseableHttpClient client = mock(CloseableHttpClient.class);
+
+  static CloseableHttpResponse response200 = mock(CloseableHttpResponse.class);
+  static CloseableHttpResponse response429 = mock(CloseableHttpResponse.class);
+  static CloseableHttpResponse response500 = mock(CloseableHttpResponse.class);
+  static HttpEntity entity200 = mock(HttpEntity.class);
+  static HttpEntity entity429 = mock(HttpEntity.class);
+  static HttpEntity entity500 = mock(HttpEntity.class);
+  static StatusLine statusLine200 = mock(StatusLine.class);
+  static StatusLine statusLine429 = mock(StatusLine.class);
+  static StatusLine statusLine500 = mock(StatusLine.class);
+  static String expectedResponseText = "expected response text.";
+
+  MlflowHttpCaller caller = new MlflowHttpCaller(new MlflowHostCredsProvider() {
+    @Override
+    public MlflowHostCreds getHostCreds() {
+      return new BasicMlflowHostCreds("http://some/host");
+    }
+
+    @Override
+    public void refresh() {
+      // pass
+    }
+  }, 4, 1, 3, client);
+
+  @BeforeSuite
+  public void beforeAll() throws IOException {
+    when(statusLine200.getStatusCode()).thenReturn(200);
+    when(response200.getStatusLine()).thenReturn(statusLine200);
+    when(entity200.getContent()).thenAnswer(
+      i -> new ByteArrayInputStream(expectedResponseText.getBytes(StandardCharsets.UTF_8))
+    );
+    when(response200.getEntity()).thenReturn(entity200);
+    when(statusLine429.getStatusCode()).thenReturn(429);
+    when(response429.getStatusLine()).thenReturn(statusLine429);
+    when(entity429.getContent()).thenAnswer(
+      i -> new ByteArrayInputStream(expectedResponseText.getBytes(StandardCharsets.UTF_8))
+    );
+    when(response429.getEntity()).thenReturn(entity429);
+    when(statusLine500.getStatusCode()).thenReturn(500);
+    when(response500.getStatusLine()).thenReturn(statusLine500);
+    when(entity500.getContent()).thenAnswer(
+      i -> new ByteArrayInputStream(expectedResponseText.getBytes(StandardCharsets.UTF_8))
+    );
+    when(response500.getEntity()).thenReturn(entity500);
+  }
+
+  @Test
+  public void testRequestsAreRetriedFor429s() throws IOException {
+    when(client.execute(any(HttpUriRequest.class)))
+            .thenReturn(response429) // sleep for 1 ms
+            .thenReturn(response429) // sleep for 2 ms
+            .thenReturn(response429) // sleep for 4 - 3 == 1ms
+            .thenReturn(response200);// last response before timing out
+    Assert.assertEquals(expectedResponseText, caller.get("some/path"));
+
+    when(client.execute(any(HttpUriRequest.class)))
+            .thenReturn(response429) // sleep for 1 ms
+            .thenReturn(response429) // sleep for 2 ms
+            .thenReturn(response429) // sleep for 4 - 3 == 1ms
+            .thenReturn(response200);// last response before timing out
+    Assert.assertEquals(expectedResponseText, caller.post("some/path", "{\"attr\":\"val\"}"));
+  }
+
+  @Test
+  public void testMaxRetryIntervalIsRespectedFor429s() throws IOException {
+    when(client.execute(any(HttpUriRequest.class)))
+            .thenReturn(response429) // sleep for 1 ms
+            .thenReturn(response429) // sleep for 2 ms
+            .thenReturn(response429) // sleep for 4 - 3 == 1ms
+            .thenReturn(response429) // last response before timing out
+            .thenReturn(response200);// should never be returned;
+    MlflowHttpException ex = Assert.expectThrows(MlflowHttpException.class,
+            () -> caller.get("some/path"));
+    Assert.assertEquals(429, ex.getStatusCode());
+
+    when(client.execute(any(HttpUriRequest.class)))
+            .thenReturn(response429) // sleep for 1 ms
+            .thenReturn(response429) // sleep for 2 ms
+            .thenReturn(response429) // sleep for 4 - 3 == 1ms
+            .thenReturn(response429) // last response before timing out
+            .thenReturn(response200);// should never be returned;
+    ex = Assert.expectThrows(MlflowHttpException.class,
+            () -> caller.post("some/path", "{\"attr\":\"val\"}"));
+    Assert.assertEquals(429, ex.getStatusCode());
+  }
+
+  @Test
+  public void testRequestsAreRetriedFor500s() throws IOException {
+    when(client.execute(any(HttpUriRequest.class)))
+            .thenReturn(response500)
+            .thenReturn(response429)
+            .thenReturn(response429)
+            .thenReturn(response500)
+            .thenReturn(response429)
+            .thenReturn(response429)
+            .thenReturn(response200);
+    Assert.assertEquals(expectedResponseText, caller.get("some/path"));
+
+    when(client.execute(any(HttpUriRequest.class)))
+            .thenReturn(response500)
+            .thenReturn(response429)
+            .thenReturn(response429)
+            .thenReturn(response500)
+            .thenReturn(response429)
+            .thenReturn(response429)
+            .thenReturn(response200);
+    Assert.assertEquals(expectedResponseText, caller.post("some/path", "{\"attr\":\"val\"}"));
+  }
+
+  @Test
+  public void testMaxRetryAttemptsIsRespectedFor500s() throws IOException {
+    when(client.execute(any(HttpUriRequest.class)))
+            .thenReturn(response500)
+            .thenReturn(response500)
+            .thenReturn(response500) // last response before timing out
+            .thenReturn(response200);// should never be returned;
+    MlflowHttpException ex = Assert.expectThrows(MlflowHttpException.class,
+            () -> caller.get("some/path"));
+    Assert.assertEquals(500, ex.getStatusCode());
+
+    when(client.execute(any(HttpUriRequest.class)))
+            .thenReturn(response500)
+            .thenReturn(response500)
+            .thenReturn(response500) // last response before timing out
+            .thenReturn(response200);// should never be returned;
+    ex = Assert.expectThrows(MlflowHttpException.class,
+            () -> caller.post("some/path", "{\"attr\":\"val\"}"));
+    Assert.assertEquals(500, ex.getStatusCode());
+  }
+}

mlflow/java/client/src/test/java/org/mlflow/tracking/HttpTest.java

@@ -0,0 +1,42 @@
+package org.mlflow.tracking;
+
+import org.testng.annotations.*;
+
+import java.io.IOException;
+
+public class HttpTest {
+  private final TestClientProvider testClientProvider = new TestClientProvider();
+
+  private MlflowClient client;
+
+  @BeforeSuite
+  public void beforeAll() throws IOException {
+    client = testClientProvider.initializeClientAndServer();
+  }
+
+  @AfterSuite
+  public void afterAll() throws InterruptedException {
+    testClientProvider.cleanupClientAndServer();
+  }
+
+  @Test(expectedExceptions = MlflowHttpException.class)
+  public void nonExistentPath() {
+    client.sendGet("BAD_PATH");
+  }
+
+  @Test(expectedExceptions = MlflowHttpException.class)   // TODO: server should throw 4xx
+  public void getExperiment_NonExistentId() {
+    client.sendGet("experiments/get?experiment_id=NON_EXISTENT_EXPERIMENT_ID");
+  }
+
+  @Test(expectedExceptions = MlflowHttpException.class) // TODO: server should throw 4xx
+  public void createExperiment_IllegalJsonSyntax() {
+    client.sendPost("experiments/create", "NOT_JSON");
+  }
+
+  @Test(expectedExceptions = MlflowHttpException.class) // TODO: server should throw 4xx
+  public void createExperiment_MissingJsonField() {
+    String data = "{\"BAD_name\": \"EXPERIMENT_NAME\"}";
+    client.sendPost("experiments/create", data);
+  }
+}

mlflow/java/client/src/test/java/org/mlflow/tracking/MlflowClientTest.java

@@ -0,0 +1,690 @@
+package org.mlflow.tracking;
+
+import java.io.File;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Optional;
+import java.util.Set;
+import java.util.Stack;
+import java.util.Vector;
+import java.util.stream.Collectors;
+import java.util.stream.LongStream;
+
+import com.google.common.collect.Lists;
+import org.apache.commons.io.FileUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.testng.Assert;
+import org.testng.annotations.AfterSuite;
+import org.testng.annotations.BeforeSuite;
+import org.testng.annotations.Test;
+
+import org.mlflow.api.proto.Service.CreateRun;
+import org.mlflow.api.proto.Service.CreateExperiment;
+import org.mlflow.api.proto.Service.Experiment;
+import org.mlflow.api.proto.Service.ExperimentTag;
+import org.mlflow.api.proto.Service.Metric;
+import org.mlflow.api.proto.Service.Param;
+import org.mlflow.api.proto.Service.Run;
+import org.mlflow.api.proto.Service.RunInfo;
+import org.mlflow.api.proto.Service.RunStatus;
+import org.mlflow.api.proto.Service.RunTag;
+import org.mlflow.api.proto.Service.ViewType;
+
+import static org.mlflow.tracking.TestUtils.assertMetric;
+import static org.mlflow.tracking.TestUtils.assertMetricHistory;
+import static org.mlflow.tracking.TestUtils.assertParam;
+import static org.mlflow.tracking.TestUtils.assertRunInfo;
+import static org.mlflow.tracking.TestUtils.assertTag;
+import static org.mlflow.tracking.TestUtils.createExperimentName;
+import static org.mlflow.tracking.TestUtils.createMetric;
+import static org.mlflow.tracking.TestUtils.createParam;
+import static org.mlflow.tracking.TestUtils.createTag;
+import static org.mlflow.tracking.TestUtils.getExperimentByName;
+
+public class MlflowClientTest {
+  private static final Logger logger = LoggerFactory.getLogger(MlflowClientTest.class);
+
+  private static double ACCURACY_SCORE = 0.9733333333333334D;
+  // NB: This can only be represented as a double (not float)
+  private static double ZERO_ONE_LOSS = 123.456789123456789D;
+  private static String MIN_SAMPLES_LEAF = "2";
+  private static String MAX_DEPTH = "3";
+  private static String USER_EMAIL = "some@email.com";
+
+  private final TestClientProvider testClientProvider = new TestClientProvider();
+  private String runId;
+
+  private MlflowClient client;
+
+  @BeforeSuite
+  public void beforeAll() throws IOException {
+    client = testClientProvider.initializeClientAndServer();
+  }
+
+  @AfterSuite
+  public void afterAll() throws InterruptedException {
+    testClientProvider.cleanupClientAndServer();
+  }
+
+  @Test
+  public void getCreateExperimentTest() {
+    String expName = createExperimentName();
+    String expId = client.createExperiment(expName);
+    Experiment exp = client.getExperiment(expId);
+    Assert.assertEquals(exp.getName(), expName);
+  }
+
+  @Test
+  public void createExperimentWithTagsTest() {
+    String expName = createExperimentName();
+    CreateExperiment.Builder request = CreateExperiment.newBuilder();
+    request.setName(expName);
+    request.addTags(ExperimentTag.newBuilder().setKey("key1").setValue("val1").build());
+    request.addTags(ExperimentTag.newBuilder().setKey("key2").setValue("val2").build());
+    String expId = client.createExperiment(request.build());
+    Experiment exp = client.getExperiment(expId);
+    Assert.assertEquals(exp.getTagsCount(), 2);
+    for (ExperimentTag tag : exp.getTagsList()) {
+      if (tag.getKey().equals("key1")) {
+        Assert.assertTrue(tag.getValue().equals("val1"));
+      }
+    }
+  }
+
+  @Test(expectedExceptions = MlflowClientException.class) // TODO: server should throw 406
+  public void createExistingExperiment() {
+    String expName = createExperimentName();
+    client.createExperiment(expName);
+    client.createExperiment(expName);
+  }
+
+  @Test
+  public void deleteAndRestoreExperiments() {
+    String expName = createExperimentName();
+    String expId = client.createExperiment(expName);
+    Assert.assertEquals(client.getExperiment(expId).getLifecycleStage(), "active");
+
+    client.deleteExperiment(expId);
+    Assert.assertEquals(client.getExperiment(expId).getLifecycleStage(), "deleted");
+
+    client.restoreExperiment(expId);
+    Assert.assertEquals(client.getExperiment(expId).getLifecycleStage(), "active");
+  }
+
+  @Test
+  public void renameExperiment() {
+    String expName = createExperimentName();
+    String newName = createExperimentName();
+
+    String expId = client.createExperiment(expName);
+    Assert.assertEquals(client.getExperiment(expId).getName(), expName);
+
+    client.renameExperiment(expId, newName);
+    Assert.assertEquals(client.getExperiment(expId).getName(), newName);
+  }
+
+  @Test
+  public void searchExperimentsTest() {
+    List<Experiment> expsBefore = client.searchExperiments().getItems();
+
+    String expName1 = createExperimentName();
+    String expId1 = client.createExperiment(expName1);
+    client.setExperimentTag(expId1, "test", "test");
+    client.setExperimentTag(expId1, "expgroup", "group1");
+
+    String expName2 = createExperimentName();
+    String expId2 = client.createExperiment(expName2);
+    client.setExperimentTag(expId2, "test", "test");
+
+    String expName3 = createExperimentName();
+    String expId3 = client.createExperiment(expName3);
+    client.setExperimentTag(expId3, "test", "test");
+    client.setExperimentTag(expId3, "expgroup", "group1");
+
+    List<Experiment> exps = client.searchExperiments().getItems();
+    Assert.assertEquals(exps.size(), 3 + expsBefore.size());
+
+    String exp1Filter = String.format("attribute.name = '%s'", expName1);
+    List<Experiment> exps1 = client.searchExperiments(exp1Filter).getItems();
+    Assert.assertEquals(exps1.size(), 1);
+    Assert.assertEquals(exps1.get(0).getExperimentId(), expId1);
+
+    String exp2Filter = String.format("attribute.name = '%s'", expName2);
+    List<Experiment> exps2 = client.searchExperiments(exp2Filter).getItems();
+    Assert.assertEquals(exps2.size(), 1);
+    Assert.assertEquals(exps2.get(0).getExperimentId(), expId2);
+
+    String expGroupFilter = String.format("tags.expgroup = 'group1'");
+    List<Experiment> expGroup = client.searchExperiments(expGroupFilter).getItems();
+    Assert.assertEquals(
+      expGroup.stream().map(exp -> exp.getExperimentId()).collect(Collectors.toSet()),
+      new HashSet<>(Arrays.asList(expId1, expId3))
+    );
+
+    client.deleteExperiment(expId2);
+
+    List<Experiment> activeExps = client.searchExperiments("").getItems();
+    Set<String> activeExpIds = activeExps.stream().map(
+        exp -> exp.getExperimentId()
+    ).collect(Collectors.toSet());
+    Assert.assertTrue(activeExpIds.contains(expId1));
+    Assert.assertTrue(activeExpIds.contains(expId3));
+    Assert.assertFalse(activeExpIds.contains(expId2));
+
+    List<Experiment> deletedExps = client.searchExperiments(
+        "", ViewType.DELETED_ONLY, 10, new ArrayList<>()
+    ).getItems();
+    Assert.assertEquals(deletedExps.size(), 1);
+    Assert.assertEquals(deletedExps.get(0).getExperimentId(), expId2);
+
+    List<String> orderedExpNames = Arrays.asList(expName1, expName2, expName3);
+    Collections.sort(orderedExpNames);
+
+    ExperimentsPage page1 = client.searchExperiments(
+      "tags.test = 'test'", ViewType.ALL, 1, Arrays.asList("attribute.name")
+    );
+    Assert.assertEquals(page1.getItems().size(), 1);
+    Assert.assertEquals(page1.getItems().get(0).getName(), orderedExpNames.get(0));
+    Assert.assertTrue(page1.getNextPageToken().isPresent());
+
+    ExperimentsPage page2 = client.searchExperiments(
+      "tags.test = 'test'",
+      ViewType.ALL,
+      2,
+      Arrays.asList("attribute.name"),
+      page1.getNextPageToken().get()
+    );
+    Assert.assertEquals(page2.getItems().size(), 2);
+    Assert.assertEquals(page2.getItems().get(0).getName(), orderedExpNames.get(1));
+    Assert.assertEquals(page2.getItems().get(1).getName(), orderedExpNames.get(2));
+    Assert.assertFalse(page2.getNextPageToken().isPresent());
+
+    ExperimentsPage nextPageFromPrevPage = (ExperimentsPage) page1.getNextPage();
+    Assert.assertEquals(nextPageFromPrevPage.getItems().size(), 1);
+    Assert.assertEquals(nextPageFromPrevPage.getItems().get(0).getName(), orderedExpNames.get(1));
+    Assert.assertTrue(nextPageFromPrevPage.getNextPageToken().isPresent());
+  }
+
+  @Test
+  public void addGetRun() {
+    // Create exp
+    String expName = createExperimentName();
+    String expId = client.createExperiment(expName);
+    logger.debug(">> TEST.0");
+
+    // Create run
+    long startTime = System.currentTimeMillis();
+
+    RunInfo runCreated = client.createRun(expId);
+    runId = runCreated.getRunUuid();
+    logger.debug("runId=" + runId);
+
+    List<RunInfo> runInfos = client.listRunInfos(expId);
+    Assert.assertEquals(runInfos.size(), 1);
+    Assert.assertEquals(runInfos.get(0).getStatus(), RunStatus.RUNNING);
+
+    // Log parameters
+    client.logParam(runId, "min_samples_leaf", MIN_SAMPLES_LEAF);
+    client.logParam(runId, "max_depth", MAX_DEPTH);
+
+    // Log metrics
+    client.logMetric(runId, "accuracy_score", ACCURACY_SCORE);
+    client.logMetric(runId, "zero_one_loss", ZERO_ONE_LOSS);
+    client.logMetric(runId, "multi_log_default_step_ts", 2.0);
+    client.logMetric(runId, "multi_log_default_step_ts", -1.0);
+    client.logMetric(runId, "multi_log_specified_step_ts", 1.0, 1000, 1);
+    client.logMetric(runId, "multi_log_specified_step_ts", 2.0, 2000, -5);
+    client.logMetric(runId, "multi_log_specified_step_ts", -3.0, 3000, 4);
+    client.logMetric(runId, "multi_log_specified_step_ts", 4.0, 2999, 4);
+
+    // Log NaNs and Infs
+    client.logMetric(runId, "nan_metric", java.lang.Double.NaN);
+    client.logMetric(runId, "pos_inf", java.lang.Double.POSITIVE_INFINITY);
+    client.logMetric(runId, "neg_inf", java.lang.Double.NEGATIVE_INFINITY);
+
+    // Log tag
+    client.setTag(runId, "user_email", USER_EMAIL);
+
+    // Update finished run
+    client.setTerminated(runId, RunStatus.FINISHED);
+    long endTime = System.currentTimeMillis();
+
+    List<RunInfo> updatedRunInfos = client.listRunInfos(expId);
+    Assert.assertEquals(updatedRunInfos.size(), 1);
+    Assert.assertEquals(updatedRunInfos.get(0).getStatus(), RunStatus.FINISHED);
+
+    // Assert run from getExperiment
+    Experiment exp = client.getExperiment(expId);
+    Assert.assertEquals(exp.getName(), expName);
+
+    // Assert run from getRun
+    Run run = client.getRun(runId);
+    RunInfo runInfo = run.getInfo();
+    assertRunInfo(runInfo, expId);
+    // verify run start and end are set in ms
+    Assert.assertTrue(runInfo.getStartTime() >= startTime);
+    Assert.assertTrue(runInfo.getEndTime() <= endTime);
+
+    // Assert parent run ID is not set.
+    Assert.assertTrue(run.getData().getTagsList().stream().noneMatch(
+            tag -> tag.getKey().equals("mlflow.parentRunId")));
+  }
+
+  @Test
+  public void setExperimentTag() {
+    // Create experiment
+    String expName = createExperimentName();
+    String expId = client.createExperiment(expName);
+    client.setExperimentTag(expId, "dataset", "imagenet1K");
+    Experiment exp = client.getExperiment(expId);
+    Assert.assertTrue(exp.getTagsCount() == 1);
+    Assert.assertTrue(exp.getTags(0).getKey().equals("dataset"));
+    Assert.assertTrue(exp.getTags(0).getValue().equals("imagenet1K"));
+    // test updating experiment tag
+    client.setExperimentTag(expId, "dataset", "birdbike");
+    exp = client.getExperiment(expId);
+    Assert.assertTrue(exp.getTagsCount() == 1);
+    Assert.assertTrue(exp.getTags(0).getKey().equals("dataset"));
+    Assert.assertTrue(exp.getTags(0).getValue().equals("birdbike"));
+    // test that setting a tag on 1 experiment does not impact another experiment.
+    String expId2 = client.createExperiment("randomExperimentName");
+    Experiment exp2 = client.getExperiment(expId2);
+    Assert.assertTrue(exp2.getTagsCount() == 0);
+    // test that setting a tag on different experiments maintain different values across experiments
+    client.setExperimentTag(expId2, "dataset", "birds200");
+    exp = client.getExperiment(expId);
+    exp2 = client.getExperiment(expId2);
+    Assert.assertTrue(exp.getTagsCount() == 1);
+    Assert.assertTrue(exp.getTags(0).getKey().equals("dataset"));
+    Assert.assertTrue(exp.getTags(0).getValue().equals("birdbike"));
+    Assert.assertTrue(exp2.getTagsCount() == 1);
+    Assert.assertTrue(exp2.getTags(0).getKey().equals("dataset"));
+    Assert.assertTrue(exp2.getTags(0).getValue().equals("birds200"));
+    // test can set multi-line tags
+    client.setExperimentTag(expId, "multiline tag", "value2\nvalue2\nvalue2");
+    exp = client.getExperiment(expId);
+    Assert.assertTrue(exp.getTagsCount() == 2);
+    for (ExperimentTag tag : exp.getTagsList()) {
+      if (tag.getKey().equals("multiline tag")) {
+        Assert.assertTrue(tag.getValue().equals("value2\nvalue2\nvalue2"));
+      }
+    }
+  }
+
+  @Test
+  public void deleteTag() {
+    // Create experiment
+    String expName = createExperimentName();
+    String expId = client.createExperiment(expName);
+
+    // Create run
+    RunInfo runCreated = client.createRun(expId);
+    String runId = runCreated.getRunUuid();
+    client.setTag(runId, "tag0", "val0");
+    client.setTag(runId, "tag1", "val1");
+    client.deleteTag(runId, "tag0");
+    Run run = client.getRun(runId);
+    // test that the tag was correctly deleted.
+    for (RunTag rt : run.getData().getTagsList()) {
+      Assert.assertTrue(!rt.getKey().equals("tag0"));
+    }
+    // test that you can't re-delete the old tag
+    try {
+      client.deleteTag(runId, "tag0");
+      Assert.fail();
+    } catch (MlflowClientException e) {
+      Assert.assertTrue(e.getMessage().contains(String.format("No tag with name: tag0 in run with id %s", runId)));
+    }
+    // test that you can't delete a tag that doesn't already exist.
+    try {
+      client.deleteTag(runId, "fakeTag");
+      Assert.fail();
+    } catch (MlflowClientException e) {
+      Assert.assertTrue(e.getMessage().contains(String.format("No tag with name: fakeTag in run with id %s", runId)));
+    }
+    // test that you can't delete a tag on a nonexistent run.
+    try {
+      client.deleteTag("fakeRunId", "fakeTag");
+      Assert.fail();
+    } catch (MlflowClientException e) {
+      Assert.assertTrue(e.getMessage().contains("Run with id=fakeRunId not found"));
+    }
+  }
+
+  @Test
+  public void searchRuns() {
+    // Create exp
+    String expName = createExperimentName();
+    String expId = client.createExperiment(expName);
+    logger.debug(">> TEST.0");
+
+    // Create run
+    String user = System.getenv("USER");
+    long startTime = System.currentTimeMillis();
+    String sourceFile = "MyFile.java";
+
+    RunInfo runCreated_1 = client.createRun(expId);
+    String runId_1 = runCreated_1.getRunUuid();
+    logger.debug("runId=" + runId_1);
+
+    RunInfo runCreated_2 = client.createRun(expId);
+    String runId_2 = runCreated_2.getRunUuid();
+    logger.debug("runId=" + runId_2);
+
+    // Log parameters
+    client.logParam(runId_1, "min_samples_leaf", MIN_SAMPLES_LEAF);
+    client.logParam(runId_2, "min_samples_leaf", MIN_SAMPLES_LEAF);
+
+    client.logParam(runId_1, "max_depth", "5");
+    client.logParam(runId_2, "max_depth", "15");
+
+    // Log metrics
+    client.logMetric(runId_1, "accuracy_score", 0.1);
+    client.logMetric(runId_1, "accuracy_score", 0.4);
+    client.logMetric(runId_2, "accuracy_score", 0.9);
+
+    // Log tag
+    client.setTag(runId_1, "user_email", USER_EMAIL);
+    client.setTag(runId_1, "test", "works");
+    client.setTag(runId_2, "test", "also works");
+
+    List<String> experimentIds = Arrays.asList(expId);
+
+    // metrics based searches
+    List<RunInfo> searchResult = client.searchRuns(experimentIds, "metrics.accuracy_score < 0");
+    Assert.assertEquals(searchResult.size(), 0);
+
+    searchResult = client.searchRuns(experimentIds, "metrics.accuracy_score > 0");
+    Assert.assertEquals(searchResult.size(), 2);
+
+    searchResult = client.searchRuns(experimentIds, "metrics.accuracy_score < 0.3");
+    Assert.assertEquals(searchResult.size(), 0);
+
+    searchResult = client.searchRuns(experimentIds, "metrics.accuracy_score < 0.5");
+    Assert.assertEquals(searchResult.get(0).getRunUuid(), runId_1);
+
+    searchResult = client.searchRuns(experimentIds, "metrics.accuracy_score > 0.5");
+    Assert.assertEquals(searchResult.get(0).getRunUuid(), runId_2);
+
+    // parameter based searches
+    searchResult = client.searchRuns(experimentIds,
+            "params.min_samples_leaf = '" + MIN_SAMPLES_LEAF + "'");
+    Assert.assertEquals(searchResult.size(), 2);
+    searchResult = client.searchRuns(experimentIds,
+            "params.min_samples_leaf != '" + MIN_SAMPLES_LEAF + "'");
+    Assert.assertEquals(searchResult.size(), 0);
+    searchResult = client.searchRuns(experimentIds, "params.max_depth = '5'");
+    Assert.assertEquals(searchResult.get(0).getRunUuid(), runId_1);
+
+    searchResult = client.searchRuns(experimentIds, "params.max_depth = '15'");
+    Assert.assertEquals(searchResult.get(0).getRunUuid(), runId_2);
+
+    // tag based search
+    searchResult = client.searchRuns(experimentIds, "tag.user_email = '" + USER_EMAIL + "'");
+    Assert.assertEquals(searchResult.get(0).getRunUuid(), runId_1);
+
+    searchResult = client.searchRuns(experimentIds, "tag.user_email != '" + USER_EMAIL + "'");
+    Assert.assertEquals(searchResult.size(), 0);
+
+    searchResult = client.searchRuns(experimentIds, "tag.test = 'works'");
+    Assert.assertEquals(searchResult.get(0).getRunUuid(), runId_1);
+
+    searchResult = client.searchRuns(experimentIds, "tag.test = 'also works'");
+    Assert.assertEquals(searchResult.get(0).getRunUuid(), runId_2);
+
+    // Paged searchRuns
+
+    List<Run> searchRuns = Lists.newArrayList(client.searchRuns(experimentIds, "", 
+            ViewType.ACTIVE_ONLY, 1000, Lists.newArrayList("metrics.accuracy_score")).getItems());
+    Assert.assertEquals(searchRuns.get(0).getInfo().getRunUuid(), runId_1);
+    Assert.assertEquals(searchRuns.get(1).getInfo().getRunUuid(), runId_2);
+
+    searchRuns = Lists.newArrayList(client.searchRuns(experimentIds, "", ViewType.ACTIVE_ONLY,
+            1000, Lists.newArrayList("params.min_samples_leaf", "metrics.accuracy_score DESC"))
+            .getItems());
+    Assert.assertEquals(searchRuns.get(1).getInfo().getRunUuid(), runId_1);
+    Assert.assertEquals(searchRuns.get(0).getInfo().getRunUuid(), runId_2);
+
+    Page<Run> page = client.searchRuns(experimentIds, "", ViewType.ACTIVE_ONLY, 1000);
+    Assert.assertEquals(page.getPageSize(), 2);
+    Assert.assertEquals(page.hasNextPage(), false);
+    Assert.assertEquals(page.getNextPageToken(), Optional.empty());
+
+    page = client.searchRuns(experimentIds, "", ViewType.ACTIVE_ONLY, 1);
+    Assert.assertEquals(page.getPageSize(), 1);
+    Assert.assertEquals(page.hasNextPage(), true);
+    Assert.assertNotEquals(page.getNextPageToken(), Optional.empty());
+
+    Page<Run> page2 = page.getNextPage();
+    Assert.assertEquals(page2.getPageSize(), 1);
+    Assert.assertEquals(page2.hasNextPage(), false);
+    Assert.assertEquals(page2.getNextPageToken(), Optional.empty());
+
+    Page<Run> page3 = page2.getNextPage();
+    Assert.assertEquals(page3.getPageSize(), 0);
+    Assert.assertEquals(page3.getNextPageToken(), Optional.empty());
+  }
+
+  @Test
+  public void createRunWithParent() {
+    String expName = createExperimentName();
+    String expId = client.createExperiment(expName);
+    RunInfo parentRun = client.createRun(expId);
+    String parentRunId = parentRun.getRunUuid();
+    RunInfo childRun = client.createRun(CreateRun.newBuilder()
+    .setExperimentId(expId)
+    .build());
+    client.setTag(childRun.getRunUuid(), "mlflow.parentRunId", parentRunId);
+    List<RunTag> childTags = client.getRun(childRun.getRunUuid()).getData().getTagsList();
+    String parentRunIdTagValue = childTags.stream()
+      .filter(t -> t.getKey().equals("mlflow.parentRunId"))
+      .findFirst()
+      .get()
+      .getValue();
+    Assert.assertEquals(parentRunIdTagValue, parentRunId);
+  }
+
+  @Test(dependsOnMethods = {"addGetRun"})
+  public void checkParamsAndMetrics() {
+
+    Run run = client.getRun(runId);
+    List<Param> params = run.getData().getParamsList();
+    Assert.assertEquals(params.size(), 2);
+    assertParam(params, "min_samples_leaf", MIN_SAMPLES_LEAF);
+    assertParam(params, "max_depth", MAX_DEPTH);
+
+    List<Metric> metrics = run.getData().getMetricsList();
+    Assert.assertEquals(metrics.size(), 7);
+    assertMetric(metrics, "accuracy_score", ACCURACY_SCORE);
+    assertMetric(metrics, "zero_one_loss", ZERO_ONE_LOSS);
+    assertMetric(metrics, "multi_log_default_step_ts", -1.0);
+    assertMetric(metrics, "multi_log_specified_step_ts", -3.0);
+    assertMetric(metrics, "nan_metric", Double.NaN);
+    assertMetric(metrics, "pos_inf", Double.POSITIVE_INFINITY);
+    assertMetric(metrics, "neg_inf", Double.NEGATIVE_INFINITY);
+    assert(metrics.get(0).getTimestamp() > 0) : metrics.get(0).getTimestamp();
+
+    List<Metric> multiDefaultMetricHistory = client.getMetricHistory(
+      runId, "multi_log_default_step_ts");
+    assertMetricHistory(multiDefaultMetricHistory, "multi_log_default_step_ts",
+      Arrays.asList(2.0, -1.0), Arrays.asList(0L, 0L));
+
+    List<Metric> multiSpecifiedMetricHistory = client.getMetricHistory(
+      runId, "multi_log_specified_step_ts");
+    assertMetricHistory(multiSpecifiedMetricHistory, "multi_log_specified_step_ts",
+      Arrays.asList(1.0, 2.0, 4.0, -3.0), Arrays.asList(1000L, 2000L, 2999L, 3000L),
+      Arrays.asList(1L, -5L, 4L, 4L));
+
+    List<RunTag> tags = run.getData().getTagsList();
+    Assert.assertEquals(tags.size(), 2);
+    assertTag(tags, "user_email", USER_EMAIL);
+  }
+
+  @Test
+  public void getMetricHistoryPagination() {
+    String expId = client.createExperiment("getMetricHistoryPagination");
+    RunInfo run = client.createRun(expId);
+    String runId = run.getRunUuid();
+    List<Long> steps = LongStream.range(0, 26000).boxed().collect(Collectors.toList());
+    for (Long step: steps) {
+      client.logMetric(runId, "random_metric", step, step, step);
+    }
+    List<Metric> metrics = client.getMetricHistory(runId, "random_metric");
+    List<Double> values = steps.stream().mapToDouble(x -> x).boxed().collect(Collectors.toList());
+    assertMetricHistory(metrics, "random_metric", values, steps);
+  }
+
+  @Test
+  public void testBatchedLogging() {
+    // Create exp
+    String expName = createExperimentName();
+    String expId = client.createExperiment(expName);
+    logger.debug(">> TEST.0");
+
+    // Test logging just metrics
+    {
+      RunInfo runCreated = client.createRun(expId);
+      String runUuid = runCreated.getRunId();
+      logger.debug("runUuid=" + runUuid);
+
+      List<Metric> metrics = new ArrayList<>(Arrays.asList(createMetric("met1", 0.081D, 10, 0),
+        createMetric("metric2", 82.3D, 100, 73), createMetric("metric3", 1.0D, 1000, 1),
+        createMetric("metric3", 2.0D, 2000, 3), createMetric("metric3", 3.0D, 0, -2)));
+      client.logBatch(runUuid, metrics, null, null);
+
+      Run run = client.getRun(runUuid);
+      Assert.assertEquals(run.getInfo().getRunId(), runUuid);
+
+      List<Metric> loggedMetrics = run.getData().getMetricsList();
+      Assert.assertEquals(loggedMetrics.size(), 3);
+      assertMetric(loggedMetrics, "met1", 0.081D, 10, 0);
+      assertMetric(loggedMetrics, "metric2", 82.3D, 100, 73);
+      assertMetric(loggedMetrics, "metric3", 2.0D, 2000, 3);
+    }
+
+    // Test logging just params
+    {
+      RunInfo runCreated = client.createRun(expId);
+      String runUuid = runCreated.getRunId();
+      logger.debug("runUuid=" + runUuid);
+
+      Set<Param> params = new HashSet<Param>(Arrays.asList(
+        createParam("p1", "this is a param string"),
+        createParam("p2", "a b"),
+        createParam("3", "x")));
+      client.logBatch(runUuid, null, params, null);
+
+      Run run = client.getRun(runUuid);
+      Assert.assertEquals(run.getInfo().getRunId(), runUuid);
+
+      List<Param> loggedParams = run.getData().getParamsList();
+      Assert.assertEquals(loggedParams.size(), 3);
+      assertParam(loggedParams, "p1", "this is a param string");
+      assertParam(loggedParams, "p2", "a b");
+      assertParam(loggedParams, "3", "x");
+    }
+
+    // Test logging just tags
+    {
+      RunInfo runCreated = client.createRun(expId);
+      String runUuid = runCreated.getRunId();
+      logger.debug("runUuid=" + runUuid);
+
+      Stack<RunTag> tags = new Stack();
+      tags.push(createTag("t1", "tagtagtag"));
+      tags.push(createTag("mlflow.runName", "myRun"));
+      client.logBatch(runUuid, null, null, tags);
+
+      Run run = client.getRun(runUuid);
+      Assert.assertEquals(run.getInfo().getRunId(), runUuid);
+
+      List<RunTag> loggedTags = run.getData().getTagsList();
+      Assert.assertEquals(loggedTags.size(), 2);
+      assertTag(loggedTags, "t1", "tagtagtag");
+      assertTag(loggedTags, "mlflow.runName", "myRun");
+    }
+
+    // All
+    {
+      RunInfo runCreated = client.createRun(expId);
+      String runUuid = runCreated.getRunId();
+      logger.debug("runUuid=" + runUuid);
+
+      List<Metric> metrics = new LinkedList<>(Arrays.asList(createMetric("m1", 32.23D, 12, 0)));
+      Vector<Param> params = new Vector<>(Arrays.asList(createParam("p1", "param1"),
+        createParam("p2", "another param")));
+      Set<RunTag> tags = new HashSet<>(Arrays.asList(createTag("t1", "t1"),
+        createTag("t2", "xx"),
+        createTag("t3", "xx"),
+        createTag("mlflow.runName", "myRun")));
+      client.logBatch(runUuid, metrics, params, tags);
+
+      Run run = client.getRun(runUuid);
+      Assert.assertEquals(run.getInfo().getRunId(), runUuid);
+
+      List<Metric> loggedMetrics = run.getData().getMetricsList();
+      Assert.assertEquals(loggedMetrics.size(), 1);
+      assertMetric(loggedMetrics, "m1", 32.23D);
+
+      List<Param> loggedParams = run.getData().getParamsList();
+      Assert.assertEquals(loggedParams.size(), 2);
+      assertParam(loggedParams, "p1", "param1");
+      assertParam(loggedParams, "p2", "another param");
+
+      List<RunTag> loggedTags = run.getData().getTagsList();
+      Assert.assertEquals(loggedTags.size(), 4);
+      assertTag(loggedTags, "t1", "t1");
+      assertTag(loggedTags, "t2", "xx");
+      assertTag(loggedTags, "t3", "xx");
+      assertTag(loggedTags, "mlflow.runName", "myRun");
+    }
+  }
+
+  @Test
+  public void deleteAndRestoreRun() {
+    String expName = createExperimentName();
+    String expId = client.createExperiment(expName);
+
+    String sourceFile = "MyFile.java";
+
+    RunInfo runCreated = client.createRun(expId);
+    Assert.assertEquals(runCreated.getLifecycleStage(), "active");
+    String deleteRunId = runCreated.getRunId();
+    client.deleteRun(deleteRunId);
+    Assert.assertEquals(client.getRun(deleteRunId).getInfo().getLifecycleStage(), "deleted");
+    client.restoreRun(deleteRunId);
+    Assert.assertEquals(client.getRun(deleteRunId).getInfo().getLifecycleStage(), "active");
+  }
+
+  @Test
+  public void testUseArtifactRepository() throws IOException {
+    String content = "Hello, Worldz!";
+
+    File tempFile = Files.createTempFile(getClass().getSimpleName(), ".txt").toFile();
+    File tempDir = Files.createTempDirectory("tempDir").toFile();
+    File tempFileForDir = Files.createTempFile(tempDir.toPath(), "file", ".txt").toFile();
+
+    FileUtils.writeStringToFile(tempFileForDir, content, StandardCharsets.UTF_8);
+    FileUtils.writeStringToFile(tempFile, content, StandardCharsets.UTF_8);
+    client.logArtifact(runId, tempFile);
+    client.logArtifact(runId, tempDir);
+
+    File downloadedArtifact = client.downloadArtifacts(runId, tempFile.getName());
+    File downloadedArtifactFromDir = client.downloadArtifacts(runId, tempDir.getName() + "/" +
+      tempFileForDir.getName());
+    String downloadedContent = FileUtils.readFileToString(downloadedArtifact,
+      StandardCharsets.UTF_8);
+    String downloadedContentFromDir = FileUtils.readFileToString(downloadedArtifactFromDir,
+      StandardCharsets.UTF_8);
+    Assert.assertEquals(content, downloadedContent);
+    Assert.assertEquals(content, downloadedContentFromDir);
+  }
+}

mlflow/java/client/src/test/java/org/mlflow/tracking/MlflowContextTest.java

@@ -0,0 +1,126 @@
+package org.mlflow.tracking;
+
+import static org.mockito.Mockito.*;
+
+import static org.mlflow.api.proto.Service.*;
+
+import org.mlflow.tracking.utils.MlflowTagConstants;
+import org.mockito.ArgumentCaptor;
+import org.testng.Assert;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.Test;
+
+import java.util.List;
+import java.util.Optional;
+
+public class MlflowContextTest {
+  private static MlflowClient mockClient;
+
+  @AfterMethod
+  public static void afterMethod() {
+    mockClient = null;
+  }
+
+  public static MlflowContext setupMlflowContext() {
+    mockClient = mock(MlflowClient.class);
+    MlflowContext mlflow = new MlflowContext(mockClient);
+    return mlflow;
+  }
+
+  @Test
+  public void testGetClient() {
+    MlflowContext mlflow = setupMlflowContext();
+    Assert.assertEquals(mlflow.getClient(), mockClient);
+  }
+
+  @Test
+  public void testSetExperimentName() {
+    // Will throw if there is no experiment with the same name.
+    {
+      MlflowContext mlflow = setupMlflowContext();
+      when(mockClient.getExperimentByName("experiment-name")).thenReturn(Optional.empty());
+      try {
+        mlflow.setExperimentName("experiment-name");
+        Assert.fail();
+      } catch (IllegalArgumentException expected) {
+      }
+    }
+
+    // Will set experiment-id if experiment is returned from getExperimentByName
+    {
+      MlflowContext mlflow = setupMlflowContext();
+      when(mockClient.getExperimentByName("experiment-name")).thenReturn(
+        Optional.of(Experiment.newBuilder().setExperimentId("123").build()));
+      mlflow.setExperimentName("experiment-name");
+      Assert.assertEquals(mlflow.getExperimentId(), "123");
+    }
+  }
+
+  @Test
+  public void testSetAndGetExperimentId() {
+      MlflowContext mlflow = setupMlflowContext();
+      mlflow.setExperimentId("apple");
+      Assert.assertEquals(mlflow.getExperimentId(), "apple");
+  }
+
+  @Test
+  public void testStartRun() {
+    // Sets the appropriate tags
+    ArgumentCaptor<CreateRun> createRunArgument = ArgumentCaptor.forClass(CreateRun.class);
+    MlflowContext mlflow = setupMlflowContext();
+    mlflow.setExperimentId("123");
+    mlflow.startRun("apple", "parent-run-id");
+    verify(mockClient).createRun(createRunArgument.capture());
+    List<RunTag> tags = createRunArgument.getValue().getTagsList();
+    Assert.assertEquals(createRunArgument.getValue().getExperimentId(), "123");
+    Assert.assertTrue(tags.contains(createRunTag(MlflowTagConstants.RUN_NAME, "apple")));
+    Assert.assertTrue(tags.contains(createRunTag(MlflowTagConstants.SOURCE_TYPE, "LOCAL")));
+    Assert.assertTrue(tags.contains(createRunTag(MlflowTagConstants.USER, System.getProperty("user.name"))));
+    Assert.assertTrue(tags.contains(createRunTag(MlflowTagConstants.PARENT_RUN_ID, "parent-run-id")));
+  }
+
+  @Test
+  public void testStartRunWithNoRunName() {
+    // Sets the appropriate tags
+    ArgumentCaptor<CreateRun> createRunArgument = ArgumentCaptor.forClass(CreateRun.class);
+    MlflowContext mlflow = setupMlflowContext();
+    mlflow.startRun();
+    verify(mockClient).createRun(createRunArgument.capture());
+    List<RunTag> tags = createRunArgument.getValue().getTagsList();
+    Assert.assertFalse(
+      tags.stream().anyMatch(tag -> tag.getKey().equals(MlflowTagConstants.RUN_NAME)));
+  }
+
+  @Test
+  public void testWithActiveRun() {
+    // Sets the appropriate tags
+    MlflowContext mlflow = setupMlflowContext();
+    mlflow.setExperimentId("123");
+    when(mockClient.createRun(any(CreateRun.class)))
+      .thenReturn(RunInfo.newBuilder().setRunId("test-id").build());
+    mlflow.withActiveRun("apple", activeRun -> {
+      Assert.assertEquals(activeRun.getId(), "test-id");
+    });
+    verify(mockClient).createRun(any(CreateRun.class));
+    verify(mockClient).setTerminated(any(), any());
+  }
+
+  @Test
+  public void testWithActiveRunNoRunName() {
+    // Sets the appropriate tags
+    MlflowContext mlflow = setupMlflowContext();
+    mlflow.setExperimentId("123");
+    when(mockClient.createRun(any(CreateRun.class)))
+      .thenReturn(RunInfo.newBuilder().setRunId("test-id").build());
+    mlflow.withActiveRun(activeRun -> {
+      Assert.assertEquals(activeRun.getId(), "test-id");
+    });
+    verify(mockClient).createRun(any(CreateRun.class));
+    verify(mockClient).setTerminated(any(), any());
+  }
+
+
+  private static RunTag createRunTag(String key, String value) {
+    return RunTag.newBuilder().setKey(key).setValue(value).build();
+  }
+}

mlflow/java/client/src/test/java/org/mlflow/tracking/MlflowProtobufMapperTest.java

@@ -0,0 +1,39 @@
+package org.mlflow.tracking;
+
+import java.lang.reflect.Type;
+import java.util.HashMap;
+import java.util.Map;
+
+import com.google.gson.Gson;
+import com.google.gson.reflect.TypeToken;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+import org.mlflow.api.proto.Service;
+
+public class MlflowProtobufMapperTest {
+  @Test
+  public void testSerializeSnakeCase() {
+    MlflowProtobufMapper mapper = new MlflowProtobufMapper();
+    String result = mapper.makeLogParam("my-id", "my-key", "my-value");
+
+    Gson gson = new Gson();
+    Type type = new TypeToken<Map<String, Object>>(){}.getType();
+    Map<String, String> serializedMessage = gson.fromJson(result, type);
+
+    Map<String, String> expectedMessage = new HashMap<>();
+    expectedMessage.put("run_uuid", "my-id");
+    expectedMessage.put("run_id", "my-id");
+    expectedMessage.put("key", "my-key");
+    expectedMessage.put("value", "my-value");
+    Assert.assertEquals(serializedMessage, expectedMessage);
+  }
+
+  @Test
+  public void testDeserializeSnakeCaseAndUnknown() {
+    MlflowProtobufMapper mapper = new MlflowProtobufMapper();
+    Service.CreateExperiment.Response result = mapper.toCreateExperimentResponse(
+      "{\"experiment_id\": 123, \"what is this field\": \"even\"}");
+    Assert.assertEquals(result.getExperimentId(), "123");
+  }
+}

mlflow/java/client/src/test/java/org/mlflow/tracking/ModelRegistryMlflowClientTest.java

@@ -0,0 +1,233 @@
+package org.mlflow.tracking;
+
+import static org.mlflow.tracking.TestUtils.createExperimentName;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doReturn;
+
+import com.google.common.collect.Lists;
+import java.io.File;
+import java.io.FilenameFilter;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.util.Arrays;
+import java.util.List;
+import java.util.UUID;
+import org.apache.commons.io.FileUtils;
+import org.mlflow.api.proto.ModelRegistry.ModelVersion;
+import org.mlflow.api.proto.ModelRegistry.RegisteredModel;
+import org.mlflow.api.proto.Service;
+import org.mlflow.api.proto.Service.RunInfo;
+import org.mockito.Mockito;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.testng.Assert;
+import org.testng.annotations.AfterTest;
+import org.testng.annotations.BeforeTest;
+import org.testng.annotations.Test;
+
+public class ModelRegistryMlflowClientTest {
+    private static final Logger logger = LoggerFactory.getLogger(ModelRegistryMlflowClientTest.class);
+
+    private static final MlflowProtobufMapper mapper = new MlflowProtobufMapper();
+
+    private final TestClientProvider testClientProvider = new TestClientProvider();
+
+    private MlflowClient client;
+    private String source;
+
+    private String modelName;
+
+    private static final String content = "Hello, Worldz!";
+
+    // As only a single `.txt` is stored as a model version artifact, this filter is used to
+    // extract the written file.
+    FilenameFilter filter = new FilenameFilter() {
+        @Override
+        public boolean accept(File f, String name) {
+            return name.endsWith(".txt");
+        }
+    };
+
+    @BeforeTest
+    public void before() throws IOException {
+        client = testClientProvider.initializeClientAndServer();
+        modelName = "Model-" + UUID.randomUUID().toString();
+
+        String expName = createExperimentName();
+        String expId = client.createExperiment(expName);
+
+        RunInfo runCreated = client.createRun(expId);
+        String runId = runCreated.getRunUuid();
+        source = String.format("runs:/%s/model", runId);
+
+        File tempDir = Files.createTempDirectory("tempDir").toFile();
+        File tempFile = Files.createTempFile(tempDir.toPath(), "file", ".txt").toFile();
+        FileUtils.writeStringToFile(tempFile, content, StandardCharsets.UTF_8);
+        client.logArtifact(runId, tempFile, "model");
+
+        client.sendPost("registered-models/create",
+                mapper.makeCreateModel(modelName));
+
+        client.sendPost("model-versions/create",
+                mapper.makeCreateModelVersion(modelName, runId, String.format("runs:/%s/model", runId)));
+    }
+
+    @AfterTest
+    public void after() throws InterruptedException {
+        testClientProvider.cleanupClientAndServer();
+    }
+
+    @Test
+    public void testGetLatestModelVersions() throws IOException {
+        // a list of stages
+        List<ModelVersion> versions = client.getLatestVersions(modelName,
+                Lists.newArrayList("None"));
+        Assert.assertEquals(versions.size(), 1);
+
+        validateDetailedModelVersion(versions.get(0), modelName, "None", "1");
+
+        client.sendPatch("model-versions/update", mapper.makeUpdateModelVersion(modelName,
+                "1"));
+        // get the latest version of all stages
+        List<ModelVersion> modelVersion = client.getLatestVersions(modelName);
+        Assert.assertEquals(modelVersion.size(), 1);
+        validateDetailedModelVersion(modelVersion.get(0), modelName, "None", "1");
+        client.sendPost("model-versions/transition-stage",
+                mapper.makeTransitionModelVersionStage(modelName, "1", "Staging"));
+        modelVersion = client.getLatestVersions(modelName);
+        Assert.assertEquals(modelVersion.size(), 1);
+        validateDetailedModelVersion(modelVersion.get(0),
+                modelName, "Staging", "1");
+    }
+
+    @Test
+    public void testGetModelVersion() {
+        ModelVersion modelVersion = client.getModelVersion(modelName, "1");
+        validateDetailedModelVersion(modelVersion, modelName, "Staging", "1");
+    }
+
+    @Test(expectedExceptions = MlflowHttpException.class, expectedExceptionsMessageRegExp = ".*RESOURCE_DOES_NOT_EXIST.*")
+    public void testGetModelVersion_NotFound() {
+        client.getModelVersion(modelName, "2");
+    }
+
+    @Test
+    public void testGetRegisteredModel() {
+	RegisteredModel model = client.getRegisteredModel(modelName);
+	Assert.assertEquals(model.getName(), modelName);
+	validateDetailedModelVersion(model.getLatestVersions(0), modelName, "Staging", "1" );
+    }
+
+    @Test
+    public void testGetModelVersionDownloadUri() {
+        String downloadUri = client.getModelVersionDownloadUri(modelName, "1");
+        Assert.assertEquals(source, downloadUri);
+    }
+
+    @Test
+    public void testDownloadModelVersion() throws IOException {
+        File tempDownloadDir = client.downloadModelVersion(modelName, "1");
+        File[] tempDownloadFile = tempDownloadDir.listFiles(filter);
+        Assert.assertEquals(tempDownloadFile.length, 1);
+        String downloadedContent = FileUtils.readFileToString(tempDownloadFile[0],
+                StandardCharsets.UTF_8);
+        Assert.assertEquals(content, downloadedContent);
+    }
+
+    @Test
+    public void testDownloadLatestModelVersion() throws IOException {
+        File tempDownloadDir = client.downloadLatestModelVersion(modelName, "None");
+        File[] tempDownloadFile = tempDownloadDir.listFiles(filter);
+        Assert.assertEquals(tempDownloadFile.length, 1);
+        String downloadedContent = FileUtils.readFileToString(tempDownloadFile[0],
+                StandardCharsets.UTF_8);
+        Assert.assertEquals(content, downloadedContent);
+    }
+
+    @Test(expectedExceptions = MlflowClientException.class)
+    public void testDownloadLatestModelVersionWhenMoreThanOneVersionIsReturned() {
+        MlflowClient mockedClient = Mockito.spy(client);
+
+        List<ModelVersion> modelVersions = Lists.newArrayList();
+        modelVersions.add(ModelVersion.newBuilder().build());
+        modelVersions.add(ModelVersion.newBuilder().build());
+        doReturn(modelVersions).when(mockedClient).getLatestVersions(any(), any());
+
+        mockedClient.downloadLatestModelVersion(modelName, "None");
+    }
+
+    private void validateDetailedModelVersion(ModelVersion details, String modelName,
+                                              String stage, String version) {
+        Assert.assertEquals(details.getCurrentStage(), stage);
+        Assert.assertEquals(details.getName(), modelName);
+        Assert.assertEquals(details.getVersion(), version);
+    }
+
+    @Test
+    public void testSearchModelVersions() {
+        List<ModelVersion> mvsBefore = client.searchModelVersions().getItems();
+
+        // create new model version of existing registered model
+        String newVersionRunId = "newVersionRunId";
+        String newVersionSource = "runs:/newVersionRunId/model";
+        client.sendPost("model-versions/create",
+                mapper.makeCreateModelVersion(modelName, newVersionRunId, newVersionSource));
+
+        // create new registered model
+        String modelName2 = "modelName2";
+        String runId2 = "runId2";
+        String source2 = "runs:/runId2/model";
+        client.sendPost("registered-models/create",
+                mapper.makeCreateModel(modelName2));
+        client.sendPost("model-versions/create",
+                mapper.makeCreateModelVersion(modelName2, runId2, source2));
+
+        List<ModelVersion> mvsAfter = client.searchModelVersions().getItems();
+        Assert.assertEquals(mvsAfter.size(), 2 + mvsBefore.size());
+
+        String filter1 = String.format("name = '%s'", modelName);
+        List<ModelVersion> mvs1 = client.searchModelVersions(filter1).getItems();
+        Assert.assertEquals(mvs1.size(), 2);
+        Assert.assertEquals(mvs1.get(0).getName(), modelName);
+        Assert.assertEquals(mvs1.get(1).getName(), modelName);
+
+        String filter2 = String.format("name = '%s'", modelName2);
+        List<ModelVersion> mvs2 = client.searchModelVersions(filter2).getItems();
+        Assert.assertEquals(mvs2.size(), 1);
+        Assert.assertEquals(mvs2.get(0).getName(), modelName2);
+        Assert.assertEquals(mvs2.get(0).getVersion(), "1");
+
+        String filter3 = String.format("run_id = '%s'", newVersionRunId);
+        List<ModelVersion> mvs3 = client.searchModelVersions(filter3).getItems();
+        Assert.assertEquals(mvs3.size(), 1);
+        Assert.assertEquals(mvs3.get(0).getName(), modelName);
+        Assert.assertEquals(mvs3.get(0).getVersion(), "2");
+
+        ModelVersionsPage page1 = client.searchModelVersions(
+            "", 1, Arrays.asList("creation_timestamp ASC")
+        );
+        Assert.assertEquals(page1.getItems().size(), 1);
+        Assert.assertEquals(page1.getItems().get(0).getName(), modelName);
+        Assert.assertTrue(page1.getNextPageToken().isPresent());
+
+        ModelVersionsPage page2 = client.searchModelVersions(
+            "",
+            2,
+            Arrays.asList("creation_timestamp ASC"),
+            page1.getNextPageToken().get()
+        );
+        Assert.assertEquals(page2.getItems().size(), 2);
+        Assert.assertEquals(page2.getItems().get(0).getName(), modelName);
+        Assert.assertEquals(page2.getItems().get(0).getRunId(), newVersionRunId);
+        Assert.assertEquals(page2.getItems().get(1).getName(), modelName2);
+        Assert.assertEquals(page2.getItems().get(1).getRunId(), runId2);
+        Assert.assertFalse(page2.getNextPageToken().isPresent());
+
+        ModelVersionsPage nextPageFromPrevPage = (ModelVersionsPage) page1.getNextPage();
+        Assert.assertEquals(nextPageFromPrevPage.getItems().size(), 1);
+        Assert.assertEquals(page2.getItems().get(0).getName(), modelName);
+        Assert.assertEquals(page2.getItems().get(0).getRunId(), newVersionRunId);
+        Assert.assertTrue(nextPageFromPrevPage.getNextPageToken().isPresent());
+    }
+}

mlflow/java/client/src/test/java/org/mlflow/tracking/TestClientProvider.java

@@ -0,0 +1,178 @@
+package org.mlflow.tracking;
+
+import java.io.*;
+import java.net.InetAddress;
+import java.net.InetSocketAddress;
+import java.net.ServerSocket;
+import java.net.Socket;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.concurrent.TimeUnit;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.mlflow.tracking.creds.MlflowHostCredsProvider;
+
+/**
+ * Provides an MLflow API client for testing. This is a real client, pointed to a real server.
+ * If the MLFLOW_TRACKING_URI environment variable is set, we will talk to the provided server;
+ * this allows running tests against existing servers. Otherwise, we will launch a local
+ * server on an ephemeral port, and manage its lifecycle.
+ */
+public class TestClientProvider {
+  private static final Logger logger = LoggerFactory.getLogger(TestClientProvider.class);
+
+  private static final long MAX_SERVER_WAIT_TIME_MILLIS = 60 * 1000;
+
+  private Process serverProcess;
+
+  private MlflowClient client;
+
+  /**
+   * Intializes an MLflow client and, if necessary, a local MLflow server process as well.
+   * Callers should always call {@link #cleanupClientAndServer()}.
+   */
+  public MlflowClient initializeClientAndServer() throws IOException {
+    if (serverProcess != null) {
+      throw new IllegalStateException("Previous server process not cleaned up");
+    }
+
+    String trackingUri = System.getenv("MLFLOW_TRACKING_URI");
+    if (trackingUri != null) {
+      logger.info("MLFLOW_TRACKING_URI was set, test will run against that server");
+      client = new MlflowClient(trackingUri);
+      return client;
+    } else {
+      Path tempDir = Files.createTempDirectory(getClass().getSimpleName());
+      String tempDBFile = tempDir.resolve("mlflow.db").toAbsolutePath().toString();
+      return startServerProcess("sqlite:///" + tempDBFile, tempDir.toString());
+    }
+  }
+
+  /**
+   * Performs any necessary cleanup on the client and server allocated by
+   * {@link #initializeClientAndServer()}. This is safe to call even if the client/server were
+   * not initialized successfully.
+   */
+  public void cleanupClientAndServer() throws InterruptedException {
+    if (client != null) {
+      client.close();
+    }
+    if (serverProcess == null) {
+      return;
+    }
+
+    try {
+      serverProcess.destroy();
+      // Do our best to ensure that the
+      boolean processTerminated = serverProcess.waitFor(30, TimeUnit.SECONDS);
+      if (!processTerminated) {
+        logger.warn("Server process did not terminate in 30 seconds, will forcibly destroy");
+        serverProcess.destroyForcibly();
+      }
+    } catch (Exception ex) {
+      logger.warn("Failed to destroy server process nicely.", ex);
+      serverProcess.destroyForcibly();
+    }
+    serverProcess = null;
+  }
+
+  public MlflowHostCredsProvider getClientHostCredsProvider(MlflowClient client) {
+    return client.getInternalHostCredsProvider();
+  }
+
+  /**
+   * Launches an "mlflow server" process locally. This requires that the "mlflow" command
+   * line client is on the local PATH (e.g., that we're within a conda environment), and that
+   * we are allowed to bind to 127.0.0.1 on ephemeral ports.
+   *
+   * Standard out and error from the server will be streamed to System.out and System.err.
+   *
+   * This method will wait until the server is up and running
+   * @param backendStoreUri the backend store uri to use
+   * @param
+   * @return MlflowClient pointed at the local server.
+   */
+  private MlflowClient startServerProcess(String backendStoreUri,
+                                          String defaultArtifactRoot) throws IOException {
+    ProcessBuilder pb = new ProcessBuilder();
+    int freePort = getFreePort();
+    String bindAddress = "127.0.0.1";
+    pb.command("mlflow", "server",
+            "--host", bindAddress,
+            "--port", "" + freePort,
+            "--backend-store-uri", backendStoreUri,
+            "--workers", "1",
+            "--default-artifact-root", defaultArtifactRoot);
+    serverProcess = pb.start();
+
+    // NB: We cannot use pb.inheritIO() because that interacts poorly with the Maven
+    // Surefire test runner (it keeps waiting for more input/output indefinitely).
+    // Therefore, we must manually drain the stdout and stderr streams.
+    drainStream(serverProcess.getInputStream(), System.out, "mlflow-server-stdout-reader");
+    drainStream(serverProcess.getErrorStream(), System.err, "mlflow-server-stderr-reader");
+
+    logger.info("Awaiting start of server on port " + freePort);
+    long startTime = System.nanoTime();
+    while (System.nanoTime() - startTime < MAX_SERVER_WAIT_TIME_MILLIS * 1000 * 1000) {
+      if (isPortOpen(bindAddress, freePort, 1)) {
+        break;
+      }
+      try {
+        Thread.sleep(1000);
+      } catch (InterruptedException e) {
+        throw new RuntimeException(e);
+      }
+    }
+    if (!isPortOpen(bindAddress, freePort, 1)) {
+      serverProcess.destroy();
+      throw new IllegalStateException("Server failed to start on port " + freePort + " after "
+        + MAX_SERVER_WAIT_TIME_MILLIS + " milliseconds.");
+    }
+
+    client = new MlflowClient("http://" + bindAddress + ":" + freePort);
+    return client;
+  }
+
+  /** Launches a new daemon Thread to drain the given InputStream into the given OutputStream. */
+  private void drainStream(InputStream inStream, PrintStream outStream, String threadName) {
+    Thread drainThread = new Thread(threadName) {
+      @Override
+      public void run() {
+        BufferedReader reader = new BufferedReader(new InputStreamReader(inStream,
+          StandardCharsets.UTF_8));
+        reader.lines().forEach(outStream::println);
+        logger.info("Drain completed on " + threadName);
+      }
+    };
+    drainThread.setDaemon(true);
+    drainThread.start();
+  }
+
+  /** Returns an ephemeral port which is very likely to be free, even in cases of contention. */
+  private int getFreePort() throws IOException {
+    // *nix systems rarely reuse recently allocated ports, so we allocate one and then release it.
+    ServerSocket sock = new ServerSocket(0);
+    int port = sock.getLocalPort();
+    sock.close();
+    return port;
+  }
+
+  /**
+   * Checks if a server is listening on the given host and port. This simply attempts to establish
+   * a TCP connection and returns false if it fails to do so within timeoutSeconds.
+   */
+  private boolean isPortOpen(String host, int port, int timeoutSeconds) {
+    try {
+      String ip = InetAddress.getByName(host).getHostAddress();
+      Socket socket = new Socket();
+      socket.connect(new InetSocketAddress(ip, port), timeoutSeconds * 1000);
+      socket.close();
+    } catch (IOException e) {
+      return false;
+    }
+    return true;
+  }
+}

mlflow/java/client/src/test/java/org/mlflow/tracking/TestUtils.java

@@ -0,0 +1,93 @@
+package org.mlflow.tracking;
+
+import java.util.*;
+
+import org.testng.Assert;
+
+import org.mlflow.api.proto.Service.*;
+
+public class TestUtils {
+
+  final static double EPSILON = 0.0001F;
+
+  static boolean equals(double a, double b) {
+    return a == b ? true : Math.abs(a - b) < EPSILON;
+  }
+
+  static void assertRunInfo(RunInfo runInfo, String experimentId) {
+    Assert.assertEquals(runInfo.getExperimentId(), experimentId);
+    Assert.assertNotEquals(runInfo.getUserId(), "");
+    Assert.assertTrue(runInfo.getStartTime() < runInfo.getEndTime());
+  }
+
+  public static void assertParam(List<Param> params, String key, String value) {
+    Assert.assertTrue(params.stream().filter(e -> e.getKey().equals(key) && e.getValue().equals(value)).findFirst().isPresent());
+  }
+
+  public static void assertMetric(List<Metric> metrics, String key, double value) {
+
+    if (Double.isNaN(value)) {
+      Assert.assertTrue(metrics.stream().filter(e -> e.getKey().equals(key) && Double.isNaN(e.getValue())).findFirst().isPresent());
+    } else if(Double.isInfinite(value) && value > 0) {
+      Assert.assertTrue(metrics.stream().filter(e -> e.getKey().equals(key) && e.getValue() >= Double.MAX_VALUE).findFirst().isPresent());
+    } else if(Double.isInfinite(value) && value < 0) {
+      Assert.assertTrue(metrics.stream().filter(e -> e.getKey().equals(key) && e.getValue() <= -Double.MAX_VALUE).findFirst().isPresent());
+    } else {
+      Assert.assertTrue(metrics.stream().filter(e -> e.getKey().equals(key) && equals(e.getValue(), value)).findFirst().isPresent());
+    }
+  }
+
+  public static void assertMetric(List<Metric> metrics, String key, double value, long timestamp, long step) {
+    Assert.assertTrue(metrics.stream().filter(
+      e -> e.getKey().equals(key) && equals(e.getValue(), value) && equals(e.getTimestamp(), timestamp)
+      && equals(e.getStep(), step)).findFirst().isPresent());
+  }
+
+  public static void assertMetricHistory(List<Metric> history, String key, List<Double> values, List<Long> steps) {
+    Assert.assertEquals(history.size(), values.size());
+    Assert.assertEquals(history.size(), steps.size());
+    for (int i = 0; i < history.size(); i++) {
+      Metric metric = history.get(i);
+      Assert.assertEquals(metric.getKey(), key);
+      Assert.assertTrue(equals(metric.getValue(), values.get(i)));
+      Assert.assertTrue(equals(metric.getStep(), steps.get(i)));
+    }
+  }
+
+  public static void assertMetricHistory(List<Metric> history, String key, List<Double> values, List<Long> timestamps, List<Long> steps) {
+    assertMetricHistory(history, key, values, steps);
+    for(int i = 0; i < history.size(); ++i) {
+      Assert.assertTrue(equals(history.get(i).getTimestamp(), timestamps.get(i)));
+    }
+  }
+
+  public static void assertTag(List<RunTag> tags, String key, String value) {
+    Assert.assertTrue(tags.stream().filter(e -> e.getKey().equals(key) && e.getValue().equals(value)).findFirst().isPresent());
+  }
+  public static java.util.Optional<Experiment> getExperimentByName(List<Experiment> exps, String expName) {
+    return exps.stream().filter(e -> e.getName().equals(expName)).findFirst();
+  }
+
+  static public String createExperimentName() {
+    return "JavaTestExp_" + UUID.randomUUID().toString();
+  }
+
+  public static Metric createMetric(String name, double value, long timestamp, long step) {
+    Metric.Builder builder = Metric.newBuilder();
+    builder.setKey(name).setValue(value).setTimestamp(timestamp);
+    builder.setKey(name).setValue(value).setStep(step);
+    return builder.build();
+  }
+
+  public static Param createParam(String name, String value) {
+    Param.Builder builder = Param.newBuilder();
+    builder.setKey(name).setValue(value);
+    return builder.build();
+  }
+
+  public static RunTag createTag(String name, String value) {
+    RunTag.Builder builder = RunTag.newBuilder();
+    builder.setKey(name).setValue(value);
+    return builder.build();
+  }
+}

mlflow/java/client/src/test/java/org/mlflow/tracking/creds/DatabricksConfigHostCredsProviderTest.java

@@ -0,0 +1,110 @@
+package org.mlflow.tracking.creds;
+
+import java.io.File;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Path;
+
+import org.apache.commons.io.FileUtils;
+import org.testng.Assert;
+import org.testng.annotations.AfterSuite;
+import org.testng.annotations.BeforeSuite;
+import org.testng.annotations.Test;
+
+public class DatabricksConfigHostCredsProviderTest {
+  private String previousUserHome = null;
+  private File databrickscfg = null;
+
+  @BeforeSuite
+  public void beforeAll() throws IOException {
+    previousUserHome = System.getProperty("user.home");
+    Path tempDir = Files.createTempDirectory(getClass().getSimpleName());
+    databrickscfg = tempDir.resolve(".databrickscfg").toFile();
+    System.setProperty("user.home", tempDir.toString());
+  }
+
+  @AfterSuite
+  public void afterAll() {
+    if (previousUserHome != null) {
+      System.setProperty("user.home", previousUserHome);
+    }
+  }
+
+  @Test
+  public void testGetTokenFromDefault() throws IOException {
+    String contents = "[DEFAULT]\nhost = https://boop.com\ntoken = dapi\n";
+    FileUtils.writeStringToFile(databrickscfg, contents, StandardCharsets.UTF_8);
+    DatabricksConfigHostCredsProvider provider = new DatabricksConfigHostCredsProvider();
+    Assert.assertEquals(provider.getHostCreds().getHost(), "https://boop.com");
+    Assert.assertEquals(provider.getHostCreds().getToken(), "dapi");
+
+    String contents2 = "[DEFAULT]\nhost = https://boop.com\ntoken=dapi2\ninsecure = TrUe";
+    FileUtils.writeStringToFile(databrickscfg, contents2, StandardCharsets.UTF_8);
+    Assert.assertEquals(provider.getHostCreds().getToken(), "dapi");
+    Assert.assertFalse(provider.getHostCreds().shouldIgnoreTlsVerification());
+    provider.refresh();
+    Assert.assertEquals(provider.getHostCreds().getToken(), "dapi2");
+    Assert.assertTrue(provider.getHostCreds().shouldIgnoreTlsVerification());
+  }
+
+  @Test
+  public void testGetUserPassFromProfile() throws IOException {
+    String contents = "[myprof]\nhost = https://boop.com\nusername = Bob\npassword = Ross\n";
+    FileUtils.writeStringToFile(databrickscfg, contents, StandardCharsets.UTF_8);
+    DatabricksConfigHostCredsProvider provider = new DatabricksConfigHostCredsProvider("myprof");
+    Assert.assertEquals(provider.getHostCreds().getHost(), "https://boop.com");
+    Assert.assertEquals(provider.getHostCreds().getUsername(), "Bob");
+    Assert.assertEquals(provider.getHostCreds().getPassword(), "Ross");
+
+    try {
+      new DatabricksConfigHostCredsProvider().getHostCreds();
+      Assert.fail();
+    } catch (IllegalStateException e) {
+      Assert.assertTrue(e.getMessage().contains("Could not find 'DEFAULT'"), e.getMessage());
+    }
+
+    try {
+      new DatabricksConfigHostCredsProvider("blah").getHostCreds();
+      Assert.fail();
+    } catch (IllegalStateException e) {
+      Assert.assertTrue(e.getMessage().contains("Could not find 'blah'"), e.getMessage());
+    }
+  }
+
+  @Test
+  public void testProfileNoHost() throws IOException {
+    String contents = "[DEFAULT]\ntoken = dabi\n";
+    FileUtils.writeStringToFile(databrickscfg, contents, StandardCharsets.UTF_8);
+    try {
+      new DatabricksConfigHostCredsProvider().getHostCreds();
+      Assert.fail();
+    } catch (IllegalStateException e) {
+      Assert.assertTrue(e.getMessage().contains("No 'host' configured"), e.getMessage());
+    }
+  }
+
+  @Test
+  public void testProfileNoAuth() throws IOException {
+    String contents = "[DEFAULT]\nhost = foo\n";
+    FileUtils.writeStringToFile(databrickscfg, contents, StandardCharsets.UTF_8);
+    try {
+      new DatabricksConfigHostCredsProvider().getHostCreds();
+      Assert.fail();
+    } catch (IllegalStateException e) {
+      Assert.assertTrue(e.getMessage().contains("No authentication configured"), e.getMessage());
+    }
+  }
+
+  @Test
+  public void testProfileNoFile() throws IOException {
+    databrickscfg.delete();
+    try {
+      new DatabricksConfigHostCredsProvider().getHostCreds();
+      Assert.fail();
+    } catch (IllegalStateException e) {
+      Assert.assertTrue(e.getMessage().contains("Could not find Databricks configuration file"),
+        e.getMessage());
+    }
+  }
+}

mlflow/java/client/src/test/java/org/mlflow/tracking/creds/DatabricksDynamicHostCredsProviderTest.java

@@ -0,0 +1,67 @@
+package org.mlflow.tracking.creds;
+
+import java.util.AbstractMap;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Set;
+
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+public class DatabricksDynamicHostCredsProviderTest {
+  private static Map<String, String> baseMap = new HashMap<>();
+  public static class MyDynamicProvider extends AbstractMap<String, String> {
+    @Override
+    public Set<Map.Entry<String, String>> entrySet() {
+      return baseMap.entrySet();
+    }
+  }
+
+  @Test
+  public void testUpdatesAfterPut() {
+    baseMap.put("host", "hello");
+    MlflowHostCredsProvider provider = DatabricksDynamicHostCredsProvider.createIfAvailable(
+      MyDynamicProvider.class.getName());
+    Assert.assertNotNull(provider);
+    Assert.assertEquals(provider.getHostCreds().getHost(), "hello");
+    Assert.assertNull(provider.getHostCreds().getToken());
+
+    baseMap.put("token", "toke");
+    Assert.assertEquals(provider.getHostCreds().getHost(), "hello");
+    Assert.assertEquals(provider.getHostCreds().getToken(), "toke");
+
+    baseMap.put("token", "toke2");
+    Assert.assertEquals(provider.getHostCreds().getHost(), "hello");
+    Assert.assertEquals(provider.getHostCreds().getToken(), "toke2");
+  }
+
+  @Test
+  public void testUsernamePassword() {
+    baseMap.put("host", "hello");
+    baseMap.put("username", "boop");
+    baseMap.put("password", "beep");
+    MlflowHostCredsProvider provider = DatabricksDynamicHostCredsProvider.createIfAvailable(
+      MyDynamicProvider.class.getName());
+    Assert.assertNotNull(provider);
+    Assert.assertEquals(provider.getHostCreds().getHost(), "hello");
+    Assert.assertEquals(provider.getHostCreds().getUsername(), "boop");
+    Assert.assertEquals(provider.getHostCreds().getPassword(), "beep");
+  }
+
+  @Test
+  public void testTlsInsecure() {
+    baseMap.put("host", "hello");
+    MlflowHostCredsProvider provider = DatabricksDynamicHostCredsProvider.createIfAvailable(
+      MyDynamicProvider.class.getName());
+    Assert.assertNotNull(provider);
+    Assert.assertEquals(provider.getHostCreds().getHost(), "hello");
+    Assert.assertFalse(provider.getHostCreds().shouldIgnoreTlsVerification());
+
+    baseMap.put("shouldIgnoreTlsVerification", "true");
+    Assert.assertTrue(provider.getHostCreds().shouldIgnoreTlsVerification());
+
+    baseMap.put("shouldIgnoreTlsVerification", "false");
+    Assert.assertFalse(provider.getHostCreds().shouldIgnoreTlsVerification());
+  }
+}
+

mlflow/java/client/src/test/java/org/mlflow/tracking/creds/HostCredsProviderChainTest.java

@@ -0,0 +1,81 @@
+package org.mlflow.tracking.creds;
+
+import org.testng.Assert;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+import org.mlflow.tracking.MlflowClientException;
+
+public class HostCredsProviderChainTest {
+  private boolean refreshCalled = false;
+  private boolean throwException = false;
+  private MlflowHostCreds providedHostCreds = null;
+  private MlflowHostCredsProvider myHostCredsProvider = new MyHostCredsProvider();
+
+  class MyHostCredsProvider implements MlflowHostCredsProvider {
+    @Override
+    public MlflowHostCreds getHostCreds() {
+      if (throwException) {
+        throw new IllegalStateException("Oh no!");
+      }
+      return providedHostCreds;
+    }
+
+    @Override
+    public void refresh() {
+      refreshCalled = true;
+    }
+  }
+
+  @BeforeMethod
+  public void beforeEach() {
+    refreshCalled = false;
+    throwException = false;
+    providedHostCreds = null;
+  }
+
+  @Test
+  public void testUseFirstIfAvailable() {
+    BasicMlflowHostCreds secondCreds = new BasicMlflowHostCreds("hosty", "tokeny");
+    HostCredsProviderChain chain = new HostCredsProviderChain(myHostCredsProvider, secondCreds);
+
+    // If we have valid credentials, we should be used.
+    providedHostCreds = new BasicMlflowHostCreds("new-host");
+    Assert.assertEquals(chain.getHostCreds().getHost(), "new-host");
+    Assert.assertNull(chain.getHostCreds().getToken());
+
+    // If our credentials are invalid, we should be skipped.
+    providedHostCreds = new BasicMlflowHostCreds(null);
+    Assert.assertEquals(chain.getHostCreds().getHost(), "hosty");
+    Assert.assertEquals(chain.getHostCreds().getToken(), "tokeny");
+
+    // If we return null, we should be skipped.
+    providedHostCreds = null;
+    Assert.assertEquals(chain.getHostCreds().getHost(), "hosty");
+    Assert.assertEquals(chain.getHostCreds().getToken(), "tokeny");
+
+    // If we return an exception, we should be skipped.
+    throwException = true;
+    Assert.assertEquals(chain.getHostCreds().getHost(), "hosty");
+    Assert.assertEquals(chain.getHostCreds().getToken(), "tokeny");
+  }
+
+  @Test
+  public void testRefreshPropagates() {
+    HostCredsProviderChain chain = new HostCredsProviderChain(myHostCredsProvider);
+    Assert.assertFalse(refreshCalled);
+    chain.refresh();
+    Assert.assertTrue(refreshCalled);
+  }
+
+  @Test
+  public void testThrowFinalError() {
+    throwException = true;
+    HostCredsProviderChain chain = new HostCredsProviderChain(myHostCredsProvider);
+    try {
+      chain.getHostCreds().getHost();
+    } catch (MlflowClientException e) {
+      Assert.assertTrue(e.getMessage().contains("Oh no!"), e.getMessage());
+    }
+  }
+}

mlflow/java/client/src/test/java/org/mlflow/tracking/utils/DatabricksContextTest.java

@@ -0,0 +1,125 @@
+package org.mlflow.tracking.utils;
+
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Maps;
+import org.testng.Assert;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+import java.util.*;
+
+public class DatabricksContextTest {
+  private static Map<String, String> baseMap = new HashMap<>();
+
+  public static class MyDynamicProvider extends AbstractMap<String, String> {
+    @Override
+    public Set<Entry<String, String>> entrySet() {
+      return baseMap.entrySet();
+    }
+  }
+
+  @BeforeMethod
+  public static void beforeMethod() {
+    baseMap = new HashMap<>();
+  }
+
+
+  @Test
+  public void testIsInDatabricksNotebook() {
+    baseMap.put("notebookId", "1");
+    DatabricksContext context = DatabricksContext.createIfAvailable(MyDynamicProvider.class.getName());
+    Assert.assertTrue(context.isInDatabricksNotebook());
+  }
+
+  @Test
+  public void testGetNotebookId() {
+    baseMap.put("notebookId", "1");
+    DatabricksContext context = DatabricksContext.createIfAvailable(MyDynamicProvider.class.getName());
+    Assert.assertEquals(context.getNotebookId(), "1");
+  }
+
+  @Test
+  public void testGetTagsWithEmptyNotebookAndJobDetails() {
+    // Will return empty map if not in Databricks notebook.
+    {
+      baseMap.put("notebookId", null);
+      baseMap.put("notebookPath", null);
+      DatabricksContext context = DatabricksContext.createIfAvailable(MyDynamicProvider.class.getName());
+      Assert.assertFalse(context.isInDatabricksNotebook());
+      Assert.assertEquals(context.getTags(), Maps.newHashMap());
+    }
+    // Will return empty map if not in Databricks job.
+    {
+      baseMap.put("jobId", null);
+      baseMap.put("jobRunId", null);
+      DatabricksContext context = DatabricksContext.createIfAvailable(MyDynamicProvider.class.getName());
+      Assert.assertFalse(context.isInDatabricksNotebook());
+      Assert.assertEquals(context.getTags(), Maps.newHashMap());
+    }
+    // Will return empty map if not config map is empty.
+    {
+      DatabricksContext context = DatabricksContext.createIfAvailable(MyDynamicProvider.class.getName());
+      Assert.assertFalse(context.isInDatabricksNotebook());
+      Assert.assertEquals(context.getTags(), Maps.newHashMap());
+    }
+  }
+
+  @Test
+  public void testGetTagsForNotebook() {
+    // Will return all tags if notebook context is set as expected.
+    {
+      baseMap = new HashMap<>();
+      Map<String, String> expectedTags = ImmutableMap.of(
+        MlflowTagConstants.DATABRICKS_NOTEBOOK_ID, "1",
+        MlflowTagConstants.DATABRICKS_NOTEBOOK_PATH, "test-path",
+        MlflowTagConstants.SOURCE_TYPE, "NOTEBOOK",
+        MlflowTagConstants.SOURCE_NAME, "test-path");
+      baseMap.put("notebookId", "1");
+      baseMap.put("notebookPath", "test-path");
+      DatabricksContext context = DatabricksContext.createIfAvailable(MyDynamicProvider.class.getName());
+      Assert.assertEquals(context.getTags(), expectedTags);
+    }
+    // Will not set notebook path tags if context doesn't have a notebookPath member.
+    {
+      baseMap = new HashMap<>();
+      Map<String, String> expectedTags = ImmutableMap.of(
+        MlflowTagConstants.DATABRICKS_NOTEBOOK_ID, "1");
+      baseMap.put("notebookId", "1");
+      baseMap.put("notebookPath", null);
+      DatabricksContext context = DatabricksContext.createIfAvailable(MyDynamicProvider.class.getName());
+      Assert.assertEquals(context.getTags(), expectedTags);
+    }
+  }
+
+  @Test
+  public void testGetTagsForJob() {
+    // Will return all tags if job context is set as expected.
+    {
+      baseMap = new HashMap<>();
+      Map<String, String> expectedTags = ImmutableMap.of(
+        MlflowTagConstants.DATABRICKS_JOB_ID, "70",
+        MlflowTagConstants.DATABRICKS_JOB_RUN_ID, "5",
+        MlflowTagConstants.DATABRICKS_JOB_TYPE, "notebook",
+        MlflowTagConstants.SOURCE_TYPE, "JOB",
+        MlflowTagConstants.SOURCE_NAME, "jobs/70/run/5");
+      baseMap.put("jobId", "70");
+      baseMap.put("jobRunId", "5");
+      baseMap.put("jobType", "notebook");
+      DatabricksContext context = DatabricksContext.createIfAvailable(MyDynamicProvider.class.getName());
+      Assert.assertEquals(context.getTags(), expectedTags);
+    }
+    // Will not set job type tag if job type is absent.
+    {
+      baseMap = new HashMap<>();
+      Map<String, String> expectedTags = ImmutableMap.of(
+        MlflowTagConstants.DATABRICKS_JOB_ID, "70",
+        MlflowTagConstants.DATABRICKS_JOB_RUN_ID, "5",
+        MlflowTagConstants.SOURCE_TYPE, "JOB",
+        MlflowTagConstants.SOURCE_NAME, "jobs/70/run/5");
+      baseMap.put("jobId", "70");
+      baseMap.put("jobRunId", "5");
+      DatabricksContext context = DatabricksContext.createIfAvailable(MyDynamicProvider.class.getName());
+      Assert.assertEquals(context.getTags(), expectedTags);
+    }
+  }
+}