chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1 @@
|
||||
dependency-reduced-pom.xml
|
||||
@@ -0,0 +1,165 @@
|
||||
# MLflow Java Client
|
||||
|
||||
Java client for [MLflow](https://mlflow.org) REST API.
|
||||
See also the MLflow [Python API](https://mlflow.org/docs/latest/python_api/index.html)
|
||||
and [REST API](https://mlflow.org/docs/latest/rest-api.html).
|
||||
|
||||
## Requirements
|
||||
|
||||
- Java 1.8
|
||||
- Maven
|
||||
- Run the [MLflow Tracking Server 0.4.2](https://mlflow.org/docs/latest/tracking.html#running-a-tracking-server)
|
||||
|
||||
## Build
|
||||
|
||||
### Build with tests
|
||||
|
||||
The MLflow Java client tests require that MLflow is on the PATH (to start a local server),
|
||||
so it is recommended to run them from within a development conda environment.
|
||||
|
||||
To build a deployable JAR and run tests:
|
||||
|
||||
```
|
||||
mvn package
|
||||
```
|
||||
|
||||
## Run
|
||||
|
||||
To run a simple sample.
|
||||
|
||||
```
|
||||
java -cp target/mlflow-java-client-0.4.2.jar \
|
||||
com.databricks.mlflow.client.samples.QuickStartDriver http://localhost:5001
|
||||
```
|
||||
|
||||
## JSON Serialization
|
||||
|
||||
MLflow Java client uses [Protobuf](https://developers.google.com/protocol-buffers/) 3.6.0 to serialize the JSON payload.
|
||||
|
||||
- [service.proto](../mlflow/protos/service.proto) - Protobuf definition of data objects.
|
||||
- [com.databricks.api.proto.mlflow.Service.java](src/main/java/com/databricks/api/proto/mlflow/Service.java) - Generated Java classes of all data objects.
|
||||
- [generate_protos.py](generate_protos.py) - One time script to generate Service.java. If service.proto changes you will need to re-run this script.
|
||||
- Javadoc can be generated by running `mvn javadoc:javadoc`. The output will be in [target/site/apidocs/index.html](target/site/apidocs/index.html).
|
||||
Here is the javadoc for [Service.java](target/site/apidocs/com/databricks/api/proto/mlflow/Service.html).
|
||||
|
||||
## Java Client API
|
||||
|
||||
See [ApiClient.java](src/main/java/org/mlflow/client/ApiClient.java)
|
||||
and [Service.java domain objects](src/main/java/org/mlflow/api/proto/mlflow/Service.java).
|
||||
|
||||
```
|
||||
Run getRun(String runId)
|
||||
RunInfo createRun()
|
||||
RunInfo createRun(String experimentId)
|
||||
RunInfo createRun(String experimentId, String appName)
|
||||
RunInfo createRun(CreateRun request)
|
||||
List<RunInfo> listRunInfos(String experimentId)
|
||||
|
||||
|
||||
List<Experiment> searchExperiments()
|
||||
GetExperiment.Response getExperiment(String experimentId)
|
||||
Optional<Experiment> getExperimentByName(String experimentName)
|
||||
long createExperiment(String experimentName)
|
||||
|
||||
void logParam(String runId, String key, String value)
|
||||
void logMetric(String runId, String key, float value)
|
||||
void setTerminated(String runId)
|
||||
void setTerminated(String runId, RunStatus status)
|
||||
void setTerminated(String runId, RunStatus status, long endTime)
|
||||
ListArtifacts.Response listArtifacts(String runId, String path)
|
||||
```
|
||||
|
||||
## Usage
|
||||
|
||||
### Java Usage
|
||||
|
||||
For a simple example see [QuickStartDriver.java](src/main/java/org/mlflow/tracking/samples/QuickStartDriver.java).
|
||||
For full examples of API coverage see the [tests](src/test/java/org/mlflow/tracking) such as [MlflowClientTest.java](src/test/java/org/mlflow/tracking/MlflowClientTest.java).
|
||||
|
||||
```
|
||||
package org.mlflow.tracking.samples;
|
||||
|
||||
import java.util.List;
|
||||
import java.util.Optional;
|
||||
|
||||
import org.apache.log4j.Level;
|
||||
import org.apache.log4j.LogManager;
|
||||
|
||||
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]);
|
||||
}
|
||||
|
||||
boolean verbose = args.length >= 2 && "true".equals(args[1]);
|
||||
if (verbose) {
|
||||
LogManager.getLogger("org.mlflow.client").setLevel(Level.DEBUG);
|
||||
}
|
||||
|
||||
System.out.println("====== createExperiment");
|
||||
String expName = "Exp_" + System.currentTimeMillis();
|
||||
String expId = client.createExperiment(expName);
|
||||
System.out.println("createExperiment: expId=" + expId);
|
||||
|
||||
System.out.println("====== getExperiment");
|
||||
GetExperiment.Response exp = client.getExperiment(expId);
|
||||
System.out.println("getExperiment: " + exp);
|
||||
|
||||
System.out.println("====== searchExperiments");
|
||||
List<Experiment> exps = client.searchExperiments();
|
||||
System.out.println("#experiments: " + exps.size());
|
||||
exps.forEach(e -> System.out.println(" Exp: " + e));
|
||||
|
||||
createRun(client, expId);
|
||||
|
||||
System.out.println("====== getExperiment again");
|
||||
GetExperiment.Response 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);
|
||||
}
|
||||
|
||||
void createRun(MlflowClient client, String expId) {
|
||||
System.out.println("====== createRun");
|
||||
|
||||
// Create run
|
||||
String sourceFile = "MyFile.java";
|
||||
RunInfo runCreated = client.createRun(expId, sourceFile);
|
||||
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);
|
||||
client.close();
|
||||
}
|
||||
}
|
||||
```
|
||||
@@ -0,0 +1,171 @@
|
||||
<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
|
||||
xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
|
||||
<modelVersion>4.0.0</modelVersion>
|
||||
|
||||
<parent>
|
||||
<groupId>org.mlflow</groupId>
|
||||
<artifactId>mlflow-parent</artifactId>
|
||||
<version>3.14.1-SNAPSHOT</version>
|
||||
<relativePath>../pom.xml</relativePath>
|
||||
</parent>
|
||||
|
||||
<artifactId>mlflow-client</artifactId>
|
||||
<packaging>jar</packaging>
|
||||
<name>MLflow Tracking API</name>
|
||||
<url>http://mlflow.org</url>
|
||||
<properties>
|
||||
<mlflow.shade.packageName>org.mlflow_project</mlflow.shade.packageName>
|
||||
</properties>
|
||||
|
||||
<dependencies>
|
||||
<dependency>
|
||||
<groupId>org.slf4j</groupId>
|
||||
<artifactId>slf4j-api</artifactId>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.apache.httpcomponents</groupId>
|
||||
<artifactId>httpclient</artifactId>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>com.google.protobuf</groupId>
|
||||
<artifactId>protobuf-java</artifactId>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>com.google.protobuf</groupId>
|
||||
<artifactId>protobuf-java-util</artifactId>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>commons-io</groupId>
|
||||
<artifactId>commons-io</artifactId>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.apache.commons</groupId>
|
||||
<artifactId>commons-configuration2</artifactId>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>io.opentelemetry.proto</groupId>
|
||||
<artifactId>opentelemetry-proto</artifactId>
|
||||
<!-- use 0.20.0 to be compatible with protobuf 3.19.6, we should upgrade later -->
|
||||
<version>0.20.0-alpha</version>
|
||||
</dependency>
|
||||
|
||||
<!-- Test dependencies -->
|
||||
<dependency>
|
||||
<groupId>org.testng</groupId>
|
||||
<artifactId>testng</artifactId>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.slf4j</groupId>
|
||||
<artifactId>slf4j-jdk14</artifactId>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.mockito</groupId>
|
||||
<artifactId>mockito-core</artifactId>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
</dependencies>
|
||||
|
||||
<build>
|
||||
<plugins>
|
||||
<plugin>
|
||||
<groupId>org.apache.maven.plugins</groupId>
|
||||
<artifactId>maven-checkstyle-plugin</artifactId>
|
||||
</plugin>
|
||||
<plugin>
|
||||
<groupId>org.apache.maven.plugins</groupId>
|
||||
<artifactId>maven-javadoc-plugin</artifactId>
|
||||
<configuration>
|
||||
<sourcepath>${project.basedir}/src/main/java</sourcepath>
|
||||
<excludePackageNames>com.databricks.api.proto.databricks:org.mlflow.scalapb_interface:org.mlflow.tracking.samples:org.mlflow.artifacts</excludePackageNames>
|
||||
<groups>
|
||||
<group>
|
||||
<title>Tracking API</title>
|
||||
<packages>org.mlflow.tracking:org.mlflow.tracking.*</packages>
|
||||
</group>
|
||||
</groups>
|
||||
</configuration>
|
||||
</plugin>
|
||||
<plugin>
|
||||
<groupId>org.apache.maven.plugins</groupId>
|
||||
<artifactId>maven-source-plugin</artifactId>
|
||||
</plugin>
|
||||
|
||||
<plugin>
|
||||
<!--
|
||||
We construct a comprehensive shaded artifact so that we do not require that downstream
|
||||
clients pick particular versions of protobuf, guava, apache http, and apache commons.
|
||||
-->
|
||||
<groupId>org.apache.maven.plugins</groupId>
|
||||
<artifactId>maven-shade-plugin</artifactId>
|
||||
<configuration>
|
||||
<minimizeJar>false</minimizeJar>
|
||||
<shadedArtifactAttached>false</shadedArtifactAttached>
|
||||
<artifactSet>
|
||||
<includes>
|
||||
<include>com.google.protobuf:protobuf-java</include>
|
||||
<include>com.google.protobuf:protobuf-java-util</include>
|
||||
<include>org.apache.httpcomponents:httpclient</include>
|
||||
<include>com.google.guava:guava</include>
|
||||
<include>com.google.code.gson:gson</include>
|
||||
<include>commons-codec:commons-codec</include>
|
||||
<include>commons-io:commons-io</include>
|
||||
<include>commons-logging:commons-logging</include>
|
||||
<include>org.apache.httpcomponents:httpcore</include>
|
||||
<include>org.apache.commons:commons-configuration2</include>
|
||||
<include>org.apache.commons:commons-lang3</include>
|
||||
<include>org.apache.commons:commons-text</include>
|
||||
<include>io.opentelemetry.proto:opentelemetry-proto</include>
|
||||
</includes>
|
||||
</artifactSet>
|
||||
<relocations>
|
||||
<relocation>
|
||||
<pattern>com.google</pattern>
|
||||
<shadedPattern>${mlflow.shade.packageName}.google</shadedPattern>
|
||||
</relocation>
|
||||
<relocation>
|
||||
<pattern>org.apache.commons</pattern>
|
||||
<shadedPattern>${mlflow.shade.packageName}.apachecommons</shadedPattern>
|
||||
</relocation>
|
||||
<relocation>
|
||||
<pattern>org.apache.http</pattern>
|
||||
<shadedPattern>${mlflow.shade.packageName}.apachehttp</shadedPattern>
|
||||
</relocation>
|
||||
<relocation>
|
||||
<!-- We have to move this package as it conflicts with other Databricks jars. -->
|
||||
<pattern>com.databricks.api.proto.databricks</pattern>
|
||||
<shadedPattern>${mlflow.shade.packageName}.databricks</shadedPattern>
|
||||
</relocation>
|
||||
<relocation>
|
||||
<pattern>io.opentelemetry</pattern>
|
||||
<shadedPattern>${mlflow.shade.packageName}.opentelemetry</shadedPattern>
|
||||
</relocation>
|
||||
</relocations>
|
||||
<transformers>
|
||||
<transformer implementation="org.apache.maven.plugins.shade.resource.DontIncludeResourceTransformer">
|
||||
<resources>
|
||||
<resource>public-suffix-list.txt</resource>
|
||||
<resource>log4j.properties</resource>
|
||||
<resource>.proto</resource>
|
||||
</resources>
|
||||
</transformer>
|
||||
</transformers>
|
||||
</configuration>
|
||||
<executions>
|
||||
<execution>
|
||||
<phase>package</phase>
|
||||
<goals>
|
||||
<goal>shade</goal>
|
||||
</goals>
|
||||
</execution>
|
||||
</executions>
|
||||
</plugin>
|
||||
</plugins>
|
||||
</build>
|
||||
<profiles>
|
||||
<profile>
|
||||
<id>do-sign</id>
|
||||
</profile>
|
||||
</profiles>
|
||||
</project>
|
||||
+8063
File diff suppressed because it is too large
Load Diff
+23539
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,2 @@
|
||||
/** Autogenerated MLflow Tracking API entity objects. */
|
||||
package org.mlflow.api.proto;
|
||||
@@ -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);
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
+314
@@ -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;
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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();
|
||||
}
|
||||
}
|
||||
@@ -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 -> {
|
||||
* 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;
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
}
|
||||
@@ -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();
|
||||
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
+103
@@ -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();
|
||||
}
|
||||
}
|
||||
+49
@@ -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
|
||||
}
|
||||
}
|
||||
+11
@@ -0,0 +1,11 @@
|
||||
package org.mlflow.tracking.creds;
|
||||
|
||||
abstract class DatabricksHostCredsProvider implements MlflowHostCredsProvider {
|
||||
|
||||
@Override
|
||||
public abstract DatabricksMlflowHostCreds getHostCreds();
|
||||
|
||||
@Override
|
||||
public abstract void refresh();
|
||||
|
||||
}
|
||||
+22
@@ -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);
|
||||
}
|
||||
}
|
||||
+48
@@ -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();
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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();
|
||||
}
|
||||
+11
@@ -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();
|
||||
}
|
||||
@@ -0,0 +1,2 @@
|
||||
/** Support for custom tracking service discovery and authentication. */
|
||||
package org.mlflow.tracking.creds;
|
||||
@@ -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;
|
||||
@@ -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();
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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";
|
||||
}
|
||||
@@ -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
|
||||
+237
@@ -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();
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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();
|
||||
}
|
||||
}
|
||||
@@ -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");
|
||||
}
|
||||
}
|
||||
+233
@@ -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());
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
@@ -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();
|
||||
}
|
||||
}
|
||||
+110
@@ -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());
|
||||
}
|
||||
}
|
||||
}
|
||||
+67
@@ -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());
|
||||
}
|
||||
}
|
||||
|
||||
+81
@@ -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());
|
||||
}
|
||||
}
|
||||
}
|
||||
+125
@@ -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);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,297 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
|
||||
<modelVersion>4.0.0</modelVersion>
|
||||
<groupId>org.mlflow</groupId>
|
||||
<artifactId>mlflow-parent</artifactId>
|
||||
<version>3.14.1-SNAPSHOT</version>
|
||||
<packaging>pom</packaging>
|
||||
<name>MLflow Parent POM</name>
|
||||
<url>http://mlflow.org</url>
|
||||
|
||||
<description>Open source platform for the machine learning lifecycle</description>
|
||||
|
||||
<repositories>
|
||||
<repository>
|
||||
<id>google-maven-central</id>
|
||||
<name>Google Maven Central</name>
|
||||
<url>https://maven-central.storage-download.googleapis.com/repos/central/data</url>
|
||||
</repository>
|
||||
</repositories>
|
||||
|
||||
<pluginRepositories>
|
||||
<pluginRepository>
|
||||
<id>google-maven-central</id>
|
||||
<name>Google Maven Central</name>
|
||||
<url>https://maven-central.storage-download.googleapis.com/repos/central/data</url>
|
||||
</pluginRepository>
|
||||
</pluginRepositories>
|
||||
|
||||
<!-- The following sections (licenses, developers, scm, and distributionManagement) are needed to
|
||||
publish pom-only packages to Maven -->
|
||||
<licenses>
|
||||
<license>
|
||||
<name>The Apache License, Version 2.0</name>
|
||||
<url>http://www.apache.org/licenses/LICENSE-2.0.txt</url>
|
||||
</license>
|
||||
</licenses>
|
||||
<developers>
|
||||
<developer>
|
||||
<name>Matei Zaharia</name>
|
||||
<email>matei@databricks.com</email>
|
||||
<organization>Databricks</organization>
|
||||
<organizationUrl>http://www.databricks.com</organizationUrl>
|
||||
</developer>
|
||||
</developers>
|
||||
<scm>
|
||||
<connection>scm:git:git://github.com/mlflow/mlflow.git</connection>
|
||||
<developerConnection>scm:git:ssh://github.com:mlflow/mlflow.git</developerConnection>
|
||||
<url>http://github.com/mlflow/mlflow/tree/master</url>
|
||||
</scm>
|
||||
|
||||
<properties>
|
||||
<mlflow-version>3.14.1-SNAPSHOT</mlflow-version>
|
||||
<maven.compiler.source>1.8</maven.compiler.source>
|
||||
<maven.compiler.target>1.8</maven.compiler.target>
|
||||
<scala.version>2.11.12</scala.version>
|
||||
<httpclient.version>4.5.6</httpclient.version>
|
||||
<git-commit.version>2.2.4</git-commit.version>
|
||||
<protobuf.version>3.25.9</protobuf.version>
|
||||
<scalapb.version>0.6.7</scalapb.version>
|
||||
<testng.version>6.14.3</testng.version>
|
||||
<slf4j.version>1.7.25</slf4j.version>
|
||||
<maven.compiler.source>1.8</maven.compiler.source>
|
||||
<maven.compiler.target>1.8</maven.compiler.target>
|
||||
</properties>
|
||||
|
||||
<modules>
|
||||
<module>client</module>
|
||||
<module>spark_2.12</module>
|
||||
<module>spark_2.13</module>
|
||||
</modules>
|
||||
|
||||
<dependencyManagement>
|
||||
<dependencies>
|
||||
<dependency>
|
||||
<groupId>junit</groupId>
|
||||
<artifactId>junit</artifactId>
|
||||
<version>4.12</version>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.mockito</groupId>
|
||||
<artifactId>mockito-core</artifactId>
|
||||
<version>2.28.2</version>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.apache.maven.plugins</groupId>
|
||||
<artifactId>maven-surefire-plugin</artifactId>
|
||||
<version>2.22.0</version>
|
||||
<scope>test</scope>
|
||||
<exclusions>
|
||||
<exclusion>
|
||||
<groupId>org.slf4j</groupId>
|
||||
<artifactId>slf4j-nop</artifactId>
|
||||
</exclusion>
|
||||
<exclusion>
|
||||
<groupId>org.slf4j</groupId>
|
||||
<artifactId>jcl-over-slf4j</artifactId>
|
||||
</exclusion>
|
||||
</exclusions>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.eclipse.jetty</groupId>
|
||||
<artifactId>jetty-server</artifactId>
|
||||
<version>9.4.11.v20180605</version>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.eclipse.jetty</groupId>
|
||||
<artifactId>jetty-servlet</artifactId>
|
||||
<version>9.4.11.v20180605</version>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.eclipse.jetty</groupId>
|
||||
<artifactId>jetty-util</artifactId>
|
||||
<version>9.4.11.v20180605</version>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>commons-io</groupId>
|
||||
<artifactId>commons-io</artifactId>
|
||||
<version>2.21.0</version>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.slf4j</groupId>
|
||||
<artifactId>slf4j-api</artifactId>
|
||||
<version>${slf4j.version}</version>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.slf4j</groupId>
|
||||
<artifactId>slf4j-simple</artifactId>
|
||||
<version>${slf4j.version}</version>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.slf4j</groupId>
|
||||
<artifactId>slf4j-jdk14</artifactId>
|
||||
<version>${slf4j.version}</version>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.apache.httpcomponents</groupId>
|
||||
<artifactId>httpclient</artifactId>
|
||||
<version>${httpclient.version}</version>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>pl.project13.maven</groupId>
|
||||
<artifactId>git-commit-id-plugin</artifactId>
|
||||
<version>${git-commit.version}</version>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.testng</groupId>
|
||||
<artifactId>testng</artifactId>
|
||||
<version>${testng.version}</version>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>com.google.protobuf</groupId>
|
||||
<artifactId>protobuf-java</artifactId>
|
||||
<version>${protobuf.version}</version>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>com.google.protobuf</groupId>
|
||||
<artifactId>protobuf-java-util</artifactId>
|
||||
<version>${protobuf.version}</version>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>javax.annotation</groupId>
|
||||
<artifactId>javax.annotation-api</artifactId>
|
||||
<version>1.2</version>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.apache.commons</groupId>
|
||||
<artifactId>commons-configuration2</artifactId>
|
||||
<version>2.14.0</version>
|
||||
</dependency>
|
||||
</dependencies>
|
||||
</dependencyManagement>
|
||||
<build>
|
||||
<plugins>
|
||||
<plugin>
|
||||
<groupId>org.sonatype.central</groupId>
|
||||
<artifactId>central-publishing-maven-plugin</artifactId>
|
||||
<version>0.8.0</version>
|
||||
<extensions>true</extensions>
|
||||
<configuration>
|
||||
<publishingServerId>central</publishingServerId>
|
||||
<centralSnapshotsUrl>https://central.sonatype.com/repository/maven-snapshots/</centralSnapshotsUrl>
|
||||
<autoPublish>true</autoPublish>
|
||||
</configuration>
|
||||
</plugin>
|
||||
</plugins>
|
||||
<pluginManagement>
|
||||
<plugins>
|
||||
<plugin>
|
||||
<groupId>org.apache.maven.plugins</groupId>
|
||||
<artifactId>maven-checkstyle-plugin</artifactId>
|
||||
<version>3.0.0</version>
|
||||
<executions>
|
||||
<execution>
|
||||
<id>checkstyle</id>
|
||||
<phase>prepare-package</phase>
|
||||
<goals>
|
||||
<goal>check</goal>
|
||||
</goals>
|
||||
<configuration>
|
||||
<failOnViolation>true</failOnViolation>
|
||||
<logViolationsToConsole>true</logViolationsToConsole>
|
||||
<checkstyleRules>
|
||||
<module name="Checker">
|
||||
<module name="TreeWalker">
|
||||
<module name="LineLength">
|
||||
<property name="max" value="100"/>
|
||||
</module>
|
||||
</module>
|
||||
</module>
|
||||
</checkstyleRules>
|
||||
<excludes>com/databricks/**,org/mlflow/api/proto/**,org/mlflow/scalapb_interface/**</excludes>
|
||||
</configuration>
|
||||
</execution>
|
||||
</executions>
|
||||
</plugin>
|
||||
<plugin>
|
||||
<groupId>org.apache.maven.plugins</groupId>
|
||||
<artifactId>maven-shade-plugin</artifactId>
|
||||
<version>3.1.0</version>
|
||||
</plugin>
|
||||
<plugin>
|
||||
<groupId>org.apache.maven.plugins</groupId>
|
||||
<artifactId>maven-javadoc-plugin</artifactId>
|
||||
<version>3.0.1</version>
|
||||
<executions>
|
||||
<execution>
|
||||
<id>attach-javadocs</id>
|
||||
<goals>
|
||||
<goal>jar</goal>
|
||||
</goals>
|
||||
</execution>
|
||||
</executions>
|
||||
</plugin>
|
||||
<plugin>
|
||||
<groupId>org.apache.maven.plugins</groupId>
|
||||
<artifactId>maven-source-plugin</artifactId>
|
||||
<version>3.0.1</version>
|
||||
<executions>
|
||||
<execution>
|
||||
<id>attach-sources</id>
|
||||
<goals>
|
||||
<goal>jar</goal>
|
||||
</goals>
|
||||
</execution>
|
||||
</executions>
|
||||
</plugin>
|
||||
<plugin>
|
||||
<groupId>net.alchim31.maven</groupId>
|
||||
<artifactId>scala-maven-plugin</artifactId>
|
||||
<version>4.2.0</version>
|
||||
</plugin>
|
||||
<plugin>
|
||||
<groupId>org.apache.maven.plugins</groupId>
|
||||
<artifactId>maven-gpg-plugin</artifactId>
|
||||
<version>1.6</version>
|
||||
</plugin>
|
||||
</plugins>
|
||||
</pluginManagement>
|
||||
</build>
|
||||
<profiles>
|
||||
<profile>
|
||||
<id>do-sign</id>
|
||||
<build>
|
||||
<plugins>
|
||||
<plugin>
|
||||
<groupId>org.apache.maven.plugins</groupId>
|
||||
<artifactId>maven-gpg-plugin</artifactId>
|
||||
<version>1.6</version>
|
||||
<executions>
|
||||
<execution>
|
||||
<id>sign-artifacts</id>
|
||||
<phase>verify</phase>
|
||||
<goals>
|
||||
<goal>sign</goal>
|
||||
</goals>
|
||||
<configuration>
|
||||
<keyname>${gpg.keyname}</keyname>
|
||||
<passphraseServerId>${gpg.keyname}</passphraseServerId>
|
||||
<gpgArguments>
|
||||
<arg>-v</arg>
|
||||
<arg>--batch</arg>
|
||||
<arg>--pinentry-mode</arg>
|
||||
<arg>loopback</arg>
|
||||
</gpgArguments>
|
||||
<executable>gpg2</executable>
|
||||
</configuration>
|
||||
</execution>
|
||||
</executions>
|
||||
</plugin>
|
||||
</plugins>
|
||||
</build>
|
||||
</profile>
|
||||
</profiles>
|
||||
</project>
|
||||
@@ -0,0 +1,15 @@
|
||||
package org.apache.spark.sql
|
||||
|
||||
import org.apache.spark.sql.execution.ui._
|
||||
import org.apache.spark.sql.execution.QueryExecution
|
||||
|
||||
|
||||
/**
|
||||
* MLflow-internal object used to access Spark-private fields in the implementation of
|
||||
* autologging Spark datasource information.
|
||||
*/
|
||||
object SparkAutologgingUtils {
|
||||
def getQueryExecution(sqlExecution: SparkListenerSQLExecutionEnd): QueryExecution = {
|
||||
sqlExecution.qe
|
||||
}
|
||||
}
|
||||
+153
@@ -0,0 +1,153 @@
|
||||
package org.mlflow.spark.autologging
|
||||
|
||||
import org.apache.spark.sql.SparkAutologgingUtils
|
||||
import org.apache.spark.sql.catalyst.plans.logical.{LeafNode, LogicalPlan}
|
||||
import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, FileTable}
|
||||
import org.apache.spark.sql.execution.datasources.{HadoopFsRelation, LogicalRelation}
|
||||
import org.apache.spark.sql.connector.catalog.Table
|
||||
import org.apache.spark.sql.execution.ui.SparkListenerSQLExecutionEnd
|
||||
import org.apache.spark.sql.sources.DataSourceRegister
|
||||
import org.slf4j.{Logger, LoggerFactory}
|
||||
|
||||
import scala.util.control.NonFatal
|
||||
|
||||
/** Case class wrapping information on a Spark datasource that was read. */
|
||||
private[autologging] case class SparkTableInfo(
|
||||
path: String,
|
||||
versionOpt: Option[String],
|
||||
formatOpt: Option[String])
|
||||
|
||||
/** Base trait for extracting Spark datasource attributes from a Spark logical plan. */
|
||||
private[autologging] trait DatasourceAttributeExtractorBase {
|
||||
protected val logger: Logger = LoggerFactory.getLogger(getClass)
|
||||
|
||||
private def getSparkTableInfoFromTable(table: Table): Option[SparkTableInfo] = {
|
||||
table match {
|
||||
case fileTable: FileTable =>
|
||||
val tableName = fileTable.name
|
||||
val splitName = tableName.split(" ")
|
||||
val lowercaseFormat = fileTable.formatName.toLowerCase()
|
||||
if (splitName.headOption.exists(head => head.toLowerCase == lowercaseFormat)) {
|
||||
Option(SparkTableInfo(splitName.tail.mkString(" "), None, Option(lowercaseFormat)))
|
||||
} else {
|
||||
Option(SparkTableInfo(fileTable.name, None, Option(fileTable.formatName)))
|
||||
}
|
||||
case other: Table =>
|
||||
Option(SparkTableInfo(other.name, None, None))
|
||||
}
|
||||
}
|
||||
|
||||
protected def maybeGetDeltaTableInfo(plan: LogicalPlan): Option[SparkTableInfo]
|
||||
|
||||
/**
|
||||
* Get SparkTableInfo representing the datasource that was read from leaf node of a Spark SQL
|
||||
* query plan
|
||||
*/
|
||||
protected def getTableInfoToLog(leafNode: LogicalPlan): Option[SparkTableInfo] = {
|
||||
val deltaInfoOpt = maybeGetDeltaTableInfo(leafNode)
|
||||
if (deltaInfoOpt.isDefined) {
|
||||
deltaInfoOpt
|
||||
} else {
|
||||
leafNode match {
|
||||
// DataSourceV2Relation was disabled in Spark 3.0.0 stable release due to some issue and
|
||||
// still not present in Spark 3.2.0. While we are not sure whether it will be back in
|
||||
// the future, we still keep this code here to support previous versions.
|
||||
case relation: DataSourceV2Relation =>
|
||||
getSparkTableInfoFromTable(relation.table)
|
||||
// This is the case for Spark 3.x except for 3.0.0-preview
|
||||
case LogicalRelation(HadoopFsRelation(index, _, _, _, fileFormat, _), _, _, _) =>
|
||||
val path: String = index.rootPaths.headOption.map(_.toString).getOrElse("unknown")
|
||||
val formatOpt = fileFormat match {
|
||||
case format: DataSourceRegister => Option(format.shortName)
|
||||
case _ => None
|
||||
}
|
||||
Option(SparkTableInfo(path, None, formatOpt))
|
||||
case _ => None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private def getLeafNodes(lp: LogicalPlan): Seq[LogicalPlan] = {
|
||||
if (lp == null) {
|
||||
return Seq.empty
|
||||
}
|
||||
if (lp.isInstanceOf[LeafNode]) {
|
||||
Seq(lp)
|
||||
} else {
|
||||
lp.children.flatMap(getLeafNodes)
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Get SparkTableInfo representing the datasource(s) that were read from a SparkListenerEvent
|
||||
* assumed to have a QueryExecution field named "qe".
|
||||
*/
|
||||
def getTableInfos(event: SparkListenerSQLExecutionEnd): Seq[SparkTableInfo] = {
|
||||
val qe = SparkAutologgingUtils.getQueryExecution(event)
|
||||
if (qe != null) {
|
||||
val leafNodes = getLeafNodes(qe.analyzed)
|
||||
leafNodes.flatMap(getTableInfoToLog)
|
||||
} else {
|
||||
Seq.empty
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/** Default datasource attribute extractor */
|
||||
object DatasourceAttributeExtractor extends DatasourceAttributeExtractorBase {
|
||||
// TODO: attempt to detect Delta table info when Delta Lake becomes compatible with Spark 3.0
|
||||
override def maybeGetDeltaTableInfo(leafNode: LogicalPlan): Option[SparkTableInfo] = None
|
||||
}
|
||||
|
||||
/** Datasource attribute extractor for REPL-ID aware environments (e.g. Databricks) */
|
||||
object ReplAwareDatasourceAttributeExtractor extends DatasourceAttributeExtractorBase {
|
||||
override protected def maybeGetDeltaTableInfo(leafNode: LogicalPlan): Option[SparkTableInfo] = {
|
||||
try {
|
||||
leafNode match {
|
||||
case lr: LogicalRelation =>
|
||||
// First, check whether LogicalRelation is a Delta table
|
||||
val obj = ReflectionUtils.getScalaObjectByName("com.databricks.sql.transaction.tahoe.DeltaTable")
|
||||
val deltaFileIndexOpt = ReflectionUtils.callMethod(obj, "unapply", Seq(lr)).asInstanceOf[Option[Any]]
|
||||
deltaFileIndexOpt.map(fileIndex => {
|
||||
val path = ReflectionUtils.getField(fileIndex, "path").toString
|
||||
val versionOpt = ReflectionUtils.maybeCallMethod(fileIndex, "tableVersion", Seq.empty).orElse(
|
||||
ReflectionUtils.maybeCallMethod(fileIndex, "version", Seq.empty)
|
||||
).map(_.toString)
|
||||
SparkTableInfo(path, versionOpt, Option("delta"))
|
||||
})
|
||||
case other => None
|
||||
}
|
||||
} catch {
|
||||
case NonFatal(e) =>
|
||||
if (logger.isTraceEnabled) {
|
||||
logger.trace(s"Unable to extract Delta table info: ${e.getMessage}")
|
||||
}
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
private def tryRedactString(value: String): String = {
|
||||
try {
|
||||
val redactor = ReflectionUtils.getScalaObjectByName(
|
||||
"com.databricks.spark.util.DatabricksSparkLogRedactor")
|
||||
ReflectionUtils.callMethod(redactor, "redact", Seq(value)).asInstanceOf[String]
|
||||
} catch {
|
||||
case NonFatal(e) =>
|
||||
if (logger.isTraceEnabled) {
|
||||
logger.trace(s"Redaction not available, using original value: ${e.getMessage}")
|
||||
}
|
||||
value
|
||||
}
|
||||
}
|
||||
|
||||
private def applyRedaction(tableInfo: SparkTableInfo): SparkTableInfo = {
|
||||
tableInfo match {
|
||||
case SparkTableInfo(path, versionOpt, formatOpt) =>
|
||||
SparkTableInfo(tryRedactString(path), versionOpt, formatOpt)
|
||||
}
|
||||
}
|
||||
|
||||
override def getTableInfos(event: SparkListenerSQLExecutionEnd): Seq[SparkTableInfo] = {
|
||||
super.getTableInfos(event).map(applyRedaction)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,45 @@
|
||||
package org.mlflow.spark.autologging
|
||||
|
||||
import java.io.{PrintWriter, StringWriter}
|
||||
|
||||
import scala.util.control.NonFatal
|
||||
|
||||
import org.slf4j.Logger
|
||||
|
||||
private[autologging] object ExceptionUtils {
|
||||
/** Helper for generating a nicely-formatted string representation of a Throwable */
|
||||
def serializeException(exc: Throwable): String = {
|
||||
val sw = new StringWriter
|
||||
exc.printStackTrace(new PrintWriter(sw))
|
||||
sw.toString
|
||||
}
|
||||
|
||||
def getUnexpectedExceptionMessage(exc: Throwable, msg: String): String = {
|
||||
s"Unexpected exception $msg. Please report this error, along with the " +
|
||||
s"following stacktrace, on https://github.com/mlflow/mlflow/issues:\n" +
|
||||
s"${ExceptionUtils.serializeException(exc)}"
|
||||
}
|
||||
|
||||
def tryAndLogSilently(logger: Logger, errorMsg: String, fn: => Any): Unit = {
|
||||
try {
|
||||
fn
|
||||
} catch {
|
||||
case NonFatal(e) =>
|
||||
if (logger.isTraceEnabled) {
|
||||
logger.trace(s"Skipping operation $errorMsg: ${e.getMessage}")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
def tryAndLogUnexpectedError(logger: Logger, errorMsg: String, fn: => Any): Unit = {
|
||||
try {
|
||||
fn
|
||||
} catch {
|
||||
case NonFatal(e) =>
|
||||
if (logger.isTraceEnabled) {
|
||||
logger.trace(getUnexpectedExceptionMessage(e, errorMsg))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
+181
@@ -0,0 +1,181 @@
|
||||
package org.mlflow.spark.autologging
|
||||
|
||||
import java.util.concurrent.{ConcurrentHashMap, ScheduledFuture, ScheduledThreadPoolExecutor, TimeUnit}
|
||||
|
||||
import py4j.Py4JException
|
||||
import org.apache.spark.scheduler.SparkListener
|
||||
|
||||
import scala.collection.JavaConverters._
|
||||
import org.apache.spark.sql.SparkSession
|
||||
import org.slf4j.LoggerFactory
|
||||
|
||||
import scala.util.{Try, Success, Failure}
|
||||
import scala.util.control.NonFatal
|
||||
|
||||
/**
|
||||
* Object exposing the actual implementation of MlflowAutologEventPublisher.
|
||||
* We opt for this pattern (an object extending a trait) so that we can mock methods of the
|
||||
* trait in testing
|
||||
*/
|
||||
object MlflowAutologEventPublisher extends MlflowAutologEventPublisherImpl {
|
||||
|
||||
}
|
||||
|
||||
/**
|
||||
* Trait implementing a publisher interface for publishing events on Spark datasource reads to
|
||||
* a set of listeners. See the design doc:
|
||||
* https://docs.google.com/document/d/11nhwZtj-rps0stxuIioFBM9lkvIh_ua45cAFy_PqdHU/edit for more
|
||||
* details.
|
||||
*/
|
||||
private[autologging] trait MlflowAutologEventPublisherImpl {
|
||||
private val logger = LoggerFactory.getLogger(getClass)
|
||||
|
||||
private[autologging] var sparkQueryListener: SparkListener = _
|
||||
private val executor = new ScheduledThreadPoolExecutor(1)
|
||||
private[autologging] val subscribers =
|
||||
new ConcurrentHashMap[String, MlflowAutologEventSubscriber]()
|
||||
private var scheduledTask: ScheduledFuture[_] = _
|
||||
|
||||
def spark: SparkSession = {
|
||||
SparkSession.getActiveSession.getOrElse(throw new RuntimeException("Unable to get active " +
|
||||
"SparkSession. Please ensure you've started a SparkSession via " +
|
||||
"SparkSession.builder.getOrCreate() before attempting to initialize Spark datasource " +
|
||||
"autologging."))
|
||||
}
|
||||
|
||||
/**
|
||||
* @returns True if Spark is running in a REPL-aware context. False otherwise.
|
||||
*/
|
||||
private def isInReplAwareContext: Boolean = {
|
||||
// Attempt to fetch the `spark.databricks.replId` property from the Spark Context.
|
||||
// The presence of this ID is a clear indication that we are in a REPL-aware environment
|
||||
val sc = spark.sparkContext
|
||||
val replId = Option(sc.getLocalProperty("spark.databricks.replId"))
|
||||
if (replId.isDefined) {
|
||||
return true
|
||||
}
|
||||
|
||||
// If the `spark.databricks.replId` is absent, we may still be in a Databricks environment,
|
||||
// which is REPL-aware. To check, we look for the presence of a Databricks-specific cluster ID
|
||||
// tag in the Spark configuration
|
||||
val clusterId = spark.conf.getOption("spark.databricks.clusterUsageTags.clusterId")
|
||||
if (clusterId.isDefined) {
|
||||
return true
|
||||
}
|
||||
|
||||
false
|
||||
}
|
||||
|
||||
// Exposed for testing
|
||||
private[autologging] def getSparkDataSourceListener: SparkListener = {
|
||||
if (isInReplAwareContext) {
|
||||
new ReplAwareSparkDataSourceListener(this)
|
||||
} else {
|
||||
new SparkDataSourceListener(this)
|
||||
}
|
||||
}
|
||||
|
||||
// Initialize Spark listener that pulls Delta query plan information & bubbles it up to registered
|
||||
// Python subscribers, along with a GC loop for removing unrespoins
|
||||
def init(gcDeadSubscribersIntervalSec: Int = 1): Unit = synchronized {
|
||||
if (sparkQueryListener == null) {
|
||||
val listener = getSparkDataSourceListener
|
||||
// NB: We take care to set the variable only after adding the Spark listener succeeds,
|
||||
// in case listener registration throws. This is defensive - adding a listener should
|
||||
// always succeed.
|
||||
spark.sparkContext.addSparkListener(listener)
|
||||
sparkQueryListener = listener
|
||||
// Schedule regular cleanup of detached subscribers, e.g. those associated with detached
|
||||
// notebooks
|
||||
val task = new Runnable {
|
||||
def run(): Unit = {
|
||||
unregisterBrokenSubscribers()
|
||||
}
|
||||
}
|
||||
scheduledTask = executor.scheduleAtFixedRate(
|
||||
task, gcDeadSubscribersIntervalSec, gcDeadSubscribersIntervalSec, TimeUnit.SECONDS)
|
||||
}
|
||||
}
|
||||
|
||||
def stop(): Unit = synchronized {
|
||||
if (sparkQueryListener != null) {
|
||||
spark.sparkContext.removeSparkListener(sparkQueryListener)
|
||||
sparkQueryListener = null
|
||||
while(!scheduledTask.cancel(false)) {
|
||||
Thread.sleep(1000)
|
||||
logger.info("Unable to cancel task for GC of unresponsive subscribers, retrying...")
|
||||
}
|
||||
subscribers.clear()
|
||||
}
|
||||
}
|
||||
|
||||
def register(subscriber: MlflowAutologEventSubscriber): Unit = synchronized {
|
||||
if (sparkQueryListener == null) {
|
||||
throw new RuntimeException("Please call init() before attempting to register a subscriber")
|
||||
}
|
||||
subscribers.put(subscriber.replId, subscriber)
|
||||
}
|
||||
|
||||
// Exposed for testing - in particular, so that we can iterate over subscribers in a specific
|
||||
// order within tests
|
||||
private[autologging] def getSubscribers: Seq[(String, MlflowAutologEventSubscriber)] = {
|
||||
subscribers.asScala.toSeq
|
||||
}
|
||||
|
||||
/** Unregister subscribers broken e.g. due to detaching of the associated Python REPL */
|
||||
private[autologging] def unregisterBrokenSubscribers(): Unit = {
|
||||
val brokenReplIds = getSubscribers.flatMap { case (replId, listener) =>
|
||||
try {
|
||||
listener.ping()
|
||||
Seq.empty
|
||||
} catch {
|
||||
case e: Py4JException =>
|
||||
logger.info(s"Subscriber with repl ID $replId not responding to health checks, " +
|
||||
s"removing it")
|
||||
Seq(replId)
|
||||
case NonFatal(e) =>
|
||||
if (logger.isTraceEnabled) {
|
||||
val msg = ExceptionUtils.getUnexpectedExceptionMessage(e, "while checking health " +
|
||||
s"of subscriber with repl ID $replId, removing it")
|
||||
logger.trace(msg)
|
||||
}
|
||||
Seq(replId)
|
||||
}
|
||||
}
|
||||
brokenReplIds.foreach { replId =>
|
||||
subscribers.remove(replId)
|
||||
}
|
||||
}
|
||||
|
||||
// https://github.com/delta-io/delta/blob/aaf3cd77dae06118f5cb7716eb2e71c791c6a148/core/src/main/scala/org/apache/spark/sql/delta/util/FileNames.scala#L26
|
||||
private val checkpointFilePattern = ".*\\d+\\.checkpoint(\\.\\d+\\.\\d+)?\\.parquet$".r.pattern
|
||||
private def isCheckpointFile(path: String): Boolean = checkpointFilePattern.matcher(path).matches()
|
||||
|
||||
private def shouldSkipPublish(path: String, format: Option[String]): Boolean = {
|
||||
// 1. Spark first loads head of the data as unknown "text" to infer the schema, which we don't want to log
|
||||
// 2. Checkpoint files don't provide useful information, so we filter them out
|
||||
(format.isEmpty || format.get == "text") || isCheckpointFile(path)
|
||||
}
|
||||
|
||||
private[autologging] def publishEvent(
|
||||
replIdOpt: Option[String],
|
||||
sparkTableInfo: SparkTableInfo): Unit = synchronized {
|
||||
sparkTableInfo match {
|
||||
case SparkTableInfo(path, version, format) if !shouldSkipPublish(path, format) =>
|
||||
for ((replId, listener) <- getSubscribers) {
|
||||
if (replIdOpt.isEmpty || replId == replIdOpt.get) {
|
||||
try {
|
||||
listener.notify(path, version.getOrElse("unknown"), format.getOrElse("unknown"))
|
||||
} catch {
|
||||
case NonFatal(e) =>
|
||||
if (logger.isTraceEnabled) {
|
||||
logger.trace(s"Unable to forward event to listener with repl ID $replId. " +
|
||||
s"Exception:\n${ExceptionUtils.serializeException(e)}")
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
case _ =>
|
||||
}
|
||||
}
|
||||
}
|
||||
+29
@@ -0,0 +1,29 @@
|
||||
package org.mlflow.spark.autologging
|
||||
|
||||
/**
|
||||
* Trait defining subscriber interface for receiving information about Spark datasource reads.
|
||||
* This trait can be implemented in Python in order to obtain datasource read
|
||||
* information, see https://www.py4j.org/advanced_topics.html#implementing-java-interfaces-from-python-callback
|
||||
*/
|
||||
trait MlflowAutologEventSubscriber {
|
||||
/**
|
||||
* Method called on datasource reads.
|
||||
* @param path Path of the datasource that was read
|
||||
* @param version Version, if applicable (e.g. for Delta tables) of datasource that was read
|
||||
* @param format Format ("csv", "json", etc) of the datasource that was read
|
||||
*/
|
||||
def notify(path: String, version: String, format: String): Any
|
||||
|
||||
/**
|
||||
* Used to verify that a subscriber is still responsive - for example,
|
||||
* in the case of a Python subscriber, invoking the ping() method from Java via a Py4J callback
|
||||
* allows us to verify that the associated Python process is still alive.
|
||||
*/
|
||||
def ping(): Unit
|
||||
|
||||
/**
|
||||
* Return the ID of the notebook associated with this subscriber, if any. The returned ID is
|
||||
* expected to be unique across all subscribers (e.g. a UUID).
|
||||
*/
|
||||
def replId: String
|
||||
}
|
||||
@@ -0,0 +1,70 @@
|
||||
package org.mlflow.spark.autologging
|
||||
|
||||
import java.lang.reflect.{Field, Method}
|
||||
|
||||
import scala.reflect.runtime.{universe => ru}
|
||||
import org.slf4j.LoggerFactory
|
||||
|
||||
import scala.collection.mutable
|
||||
|
||||
private[autologging] object ReflectionUtils {
|
||||
private val logger = LoggerFactory.getLogger(getClass)
|
||||
private val rm = ru.runtimeMirror(getClass.getClassLoader)
|
||||
|
||||
/** Get Scala object by its fully-qualified name */
|
||||
def getScalaObjectByName(name: String): Any = {
|
||||
val module = rm.staticModule(name)
|
||||
val obj = rm.reflectModule(module)
|
||||
obj.instance
|
||||
}
|
||||
|
||||
def getField(obj: Any, fieldName: String): Any = {
|
||||
var declaredFields: mutable.Buffer[Field] = obj.getClass.getDeclaredFields.toBuffer
|
||||
var superClass = obj.getClass.getSuperclass
|
||||
while (superClass != null) {
|
||||
declaredFields = declaredFields ++ superClass.getDeclaredFields
|
||||
superClass = superClass.getSuperclass
|
||||
}
|
||||
val field = declaredFields.find(_.getName == fieldName).getOrElse {
|
||||
throw new RuntimeException(s"Unable to get field '$fieldName' in object with class " +
|
||||
s"${obj.getClass.getName}. Available fields: " +
|
||||
s"[${declaredFields.map(_.getName).mkString(", ")}]")
|
||||
}
|
||||
field.setAccessible(true)
|
||||
field.get(obj)
|
||||
}
|
||||
|
||||
/**
|
||||
* Call method with provided name on the specified object. The method name is assumed to be
|
||||
* unique
|
||||
*/
|
||||
def callMethod(obj: Any, name: Any, args: Seq[Object]): Any = {
|
||||
var declaredMethods: mutable.Buffer[Method] = obj.getClass.getDeclaredMethods.toBuffer
|
||||
var superClass = obj.getClass.getSuperclass
|
||||
while (superClass != null) {
|
||||
declaredMethods = declaredMethods ++ superClass.getDeclaredMethods
|
||||
superClass = superClass.getSuperclass
|
||||
}
|
||||
val method = declaredMethods.find(_.getName == name).getOrElse(
|
||||
throw new RuntimeException(s"Unable to find method with name $name of object with class " +
|
||||
s"${obj.getClass.getName}. Available methods: " +
|
||||
s"[${declaredMethods.map(_.getName).mkString(", ")}]"))
|
||||
method.invoke(obj, args: _*)
|
||||
}
|
||||
|
||||
def maybeCallMethod(obj: Any, name: Any, args: Seq[Object]): Option[Any] = {
|
||||
var declaredMethods: mutable.Buffer[Method] = obj.getClass.getDeclaredMethods.toBuffer
|
||||
var superClass = obj.getClass.getSuperclass
|
||||
while (superClass != null) {
|
||||
declaredMethods = declaredMethods ++ superClass.getDeclaredMethods
|
||||
superClass = superClass.getSuperclass
|
||||
}
|
||||
|
||||
val methodOpt = declaredMethods.find(_.getName == name)
|
||||
|
||||
methodOpt match {
|
||||
case Some(method) => Some(method.invoke(obj, args: _*))
|
||||
case None => None
|
||||
}
|
||||
}
|
||||
}
|
||||
+55
@@ -0,0 +1,55 @@
|
||||
package org.mlflow.spark.autologging
|
||||
|
||||
import org.apache.spark.scheduler._
|
||||
import org.apache.spark.sql.catalyst.plans.logical.{LeafNode, LogicalPlan}
|
||||
import org.apache.spark.sql.execution.ui.{SparkListenerSQLExecutionEnd, SparkListenerSQLExecutionStart}
|
||||
import org.apache.spark.sql.execution.{QueryExecution, SQLExecution}
|
||||
|
||||
import scala.collection.JavaConverters._
|
||||
import scala.collection.mutable
|
||||
|
||||
/**
|
||||
* Implementation of the SparkListener interface used to detect Spark datasource reads.
|
||||
* and notify subscribers. Used in REPL-ID aware environments (e.g. Databricks)
|
||||
*/
|
||||
class ReplAwareSparkDataSourceListener(
|
||||
publisher: MlflowAutologEventPublisherImpl = MlflowAutologEventPublisher)
|
||||
extends SparkDataSourceListener(publisher) {
|
||||
private val executionIdToReplId = mutable.Map[Long, String]()
|
||||
|
||||
override protected def getDatasourceAttributeExtractor: DatasourceAttributeExtractorBase = {
|
||||
ReplAwareDatasourceAttributeExtractor
|
||||
}
|
||||
|
||||
private[autologging] def getProperties(event: SparkListenerJobStart): Map[String, String] = {
|
||||
Option(event.properties).map(_.asScala.toMap).getOrElse(Map.empty)
|
||||
}
|
||||
|
||||
override def onJobStart(event: SparkListenerJobStart): Unit = {
|
||||
val properties = getProperties(event)
|
||||
val executionIdOpt = properties.get(SQLExecution.EXECUTION_ID_KEY).map(_.toLong)
|
||||
val replIdOpt = properties.get("spark.databricks.replId")
|
||||
|
||||
(executionIdOpt, replIdOpt) match {
|
||||
case (Some(executionId), Some(replId)) =>
|
||||
executionIdToReplId.put(executionId, replId)
|
||||
case _ =>
|
||||
logger.trace(s"Skipping datasource autolog - required properties not available")
|
||||
}
|
||||
}
|
||||
|
||||
protected[autologging] override def onSQLExecutionEnd(event: SparkListenerSQLExecutionEnd): Unit = {
|
||||
val extractor = getDatasourceAttributeExtractor
|
||||
val tableInfos = extractor.getTableInfos(event)
|
||||
val replIdOpt = popReplIdOpt(event)
|
||||
if (replIdOpt.isDefined) {
|
||||
tableInfos.foreach { tableInfo =>
|
||||
publisher.publishEvent(replIdOpt = replIdOpt, sparkTableInfo = tableInfo)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private def popReplIdOpt(event: SparkListenerSQLExecutionEnd): Option[String] = {
|
||||
executionIdToReplId.remove(event.executionId)
|
||||
}
|
||||
}
|
||||
+41
@@ -0,0 +1,41 @@
|
||||
package org.mlflow.spark.autologging
|
||||
|
||||
import org.apache.spark.scheduler._
|
||||
import org.apache.spark.sql.execution.ui.SparkListenerSQLExecutionEnd
|
||||
import org.slf4j.LoggerFactory
|
||||
import scala.util.control.NonFatal
|
||||
|
||||
|
||||
/**
|
||||
* Implementation of the SparkListener interface used to detect Spark datasource reads.
|
||||
* and notify subscribers.
|
||||
*/
|
||||
class SparkDataSourceListener(
|
||||
publisher: MlflowAutologEventPublisherImpl = MlflowAutologEventPublisher) extends SparkListener {
|
||||
protected val logger = LoggerFactory.getLogger(getClass)
|
||||
|
||||
protected def getDatasourceAttributeExtractor: DatasourceAttributeExtractorBase = {
|
||||
DatasourceAttributeExtractor
|
||||
}
|
||||
|
||||
protected[autologging] def onSQLExecutionEnd(event: SparkListenerSQLExecutionEnd): Unit = {
|
||||
val extractor = getDatasourceAttributeExtractor
|
||||
val tableInfos = extractor.getTableInfos(event)
|
||||
tableInfos.foreach { tableInfo =>
|
||||
publisher.publishEvent(replIdOpt = None, sparkTableInfo = tableInfo)
|
||||
}
|
||||
}
|
||||
|
||||
override def onOtherEvent(event: SparkListenerEvent): Unit = {
|
||||
event match {
|
||||
case e: SparkListenerSQLExecutionEnd =>
|
||||
try {
|
||||
onSQLExecutionEnd(e)
|
||||
} catch {
|
||||
case NonFatal(ex) =>
|
||||
logger.trace(s"Skipping datasource autolog: ${ex.getMessage}")
|
||||
}
|
||||
case _ =>
|
||||
}
|
||||
}
|
||||
}
|
||||
+12
@@ -0,0 +1,12 @@
|
||||
package org.apache.spark.mlflow
|
||||
|
||||
import org.apache.spark.scheduler.SparkListenerInterface
|
||||
import org.apache.spark.sql.SparkSession
|
||||
import org.mlflow.spark.autologging.SparkDataSourceListener
|
||||
|
||||
/** Test-only object used to access Spark-private fields */
|
||||
object MlflowSparkAutologgingTestUtils {
|
||||
def getListeners(spark: SparkSession): Seq[SparkListenerInterface] = {
|
||||
spark.sparkContext.listenerBus.findListenersByClass[SparkDataSourceListener]
|
||||
}
|
||||
}
|
||||
+82
@@ -0,0 +1,82 @@
|
||||
package org.mlflow.spark.autologging
|
||||
|
||||
import org.scalatest.funsuite.AnyFunSuite
|
||||
|
||||
object TestObject {
|
||||
def myMethod: String = "hi"
|
||||
}
|
||||
|
||||
object TestFileIndex {
|
||||
def version: String = "1.0"
|
||||
}
|
||||
|
||||
abstract class TestAbstractClass {
|
||||
protected def addNumbers(x: Int, y: Int): Int = x + y
|
||||
protected val myProtectedVal: Int = 5
|
||||
}
|
||||
|
||||
class RealClass extends TestAbstractClass {
|
||||
private val myField: String = "myCoolVal"
|
||||
def subclassMethod(x: Int): Int = x * x
|
||||
}
|
||||
|
||||
class ReflectionUtilsSuite extends AnyFunSuite {
|
||||
|
||||
test("Can get private & protected fields of an object via reflection") {
|
||||
val obj = new RealClass()
|
||||
val field0 = ReflectionUtils.getField(obj, "myField").asInstanceOf[String]
|
||||
assert(field0 == "myCoolVal")
|
||||
val field1 = ReflectionUtils.getField(obj, "myProtectedVal").asInstanceOf[Int]
|
||||
assert(field1 == 5)
|
||||
}
|
||||
|
||||
test("Can call methods via reflection") {
|
||||
val obj = new RealClass()
|
||||
val args0: Seq[Object] = Seq[Integer](3)
|
||||
val res0 = ReflectionUtils.callMethod(obj, "subclassMethod", args0).asInstanceOf[Int]
|
||||
assert(res0 == 9)
|
||||
val args1: Seq[Object] = Seq[Integer](5, 6)
|
||||
val res1 = ReflectionUtils.callMethod(obj, "addNumbers", args1).asInstanceOf[Int]
|
||||
assert(res1 == 11)
|
||||
}
|
||||
|
||||
test("Can get Scala object and call methods via reflection") {
|
||||
val obj = ReflectionUtils.getScalaObjectByName("org.mlflow.spark.autologging.TestObject")
|
||||
val res = ReflectionUtils.callMethod(obj, "myMethod", Seq.empty).asInstanceOf[String]
|
||||
assert(res == "hi")
|
||||
}
|
||||
|
||||
test("maybeCallMethod None if method not found") {
|
||||
val obj = new RealClass()
|
||||
val res = ReflectionUtils.maybeCallMethod(obj, "nonExistentMethod", Seq.empty)
|
||||
|
||||
assert (res.isEmpty)
|
||||
}
|
||||
|
||||
test("maybeCallMethod invokes the method if the method is found") {
|
||||
val obj = ReflectionUtils.getScalaObjectByName("org.mlflow.spark.autologging.TestObject")
|
||||
val res0 = ReflectionUtils.maybeCallMethod(obj, "myMethod", Seq.empty).getOrElse("")
|
||||
assert(res0 == "hi")
|
||||
}
|
||||
|
||||
test("chaining maybeCallMethod works") {
|
||||
val fileIndex = ReflectionUtils.getScalaObjectByName("org.mlflow.spark.autologging.TestFileIndex")
|
||||
|
||||
val versionOpt0 = ReflectionUtils.maybeCallMethod(fileIndex, "version", Seq.empty).orElse(
|
||||
Option("second thing")
|
||||
).map(_.toString)
|
||||
assert(versionOpt0 == Some("1.0"))
|
||||
|
||||
// if only the second method exists, return it
|
||||
val versionOpt1 = ReflectionUtils.maybeCallMethod(fileIndex, "tableVersion", Seq.empty).orElse(
|
||||
ReflectionUtils.maybeCallMethod(fileIndex, "version", Seq.empty)
|
||||
).map(_.toString)
|
||||
assert(versionOpt1 == Some("1.0"))
|
||||
|
||||
// if both don't exist, just return None
|
||||
val versionOpt2 = ReflectionUtils.maybeCallMethod(fileIndex, "tableVersion", Seq.empty).orElse(
|
||||
ReflectionUtils.maybeCallMethod(fileIndex, "anotherTableVersion", Seq.empty)
|
||||
).map(_.toString)
|
||||
assert(versionOpt2 == None)
|
||||
}
|
||||
}
|
||||
+444
@@ -0,0 +1,444 @@
|
||||
package org.mlflow.spark.autologging
|
||||
|
||||
import java.io.File
|
||||
import java.nio.file.{Files, Path, Paths}
|
||||
import java.util.UUID
|
||||
|
||||
import org.apache.spark.mlflow.MlflowSparkAutologgingTestUtils
|
||||
import org.apache.spark.sql.execution.ui.SparkListenerSQLExecutionEnd
|
||||
import org.apache.spark.sql.{Row, SparkSession}
|
||||
import org.apache.spark.sql.types.{IntegerType, StringType, StructField, StructType}
|
||||
import org.mockito.Matchers.{any, eq => meq}
|
||||
import org.mockito.Mockito._
|
||||
import org.scalatest.{BeforeAndAfterAll, BeforeAndAfterEach}
|
||||
import org.scalatest.funsuite.AnyFunSuite
|
||||
import org.scalatest.matchers.should.Matchers
|
||||
|
||||
import scala.collection.mutable.ArrayBuffer
|
||||
|
||||
private[autologging] class MockSubscriber extends MlflowAutologEventSubscriber {
|
||||
private val uuid: String = UUID.randomUUID().toString
|
||||
override def replId: String = {
|
||||
uuid
|
||||
}
|
||||
|
||||
override def notify(path: String, version: String, format: String): Unit = {
|
||||
}
|
||||
|
||||
override def ping(): Unit = {}
|
||||
}
|
||||
|
||||
private[autologging] class BrokenSubscriber extends MockSubscriber {
|
||||
override def ping(): Unit = {
|
||||
throw new RuntimeException("Oh no, failing ping!")
|
||||
}
|
||||
|
||||
override def notify(path: String, version: String, format: String): Unit = {
|
||||
throw new RuntimeException("Unable to notify subscriber!")
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
class SparkAutologgingSuite extends AnyFunSuite with Matchers with BeforeAndAfterAll
|
||||
with BeforeAndAfterEach {
|
||||
|
||||
var spark: SparkSession = getOrCreateSparkSession()
|
||||
|
||||
var tempDir: Path = _
|
||||
var formatToTablePath: Map[String, String] = _
|
||||
var deltaTablePath: String = _
|
||||
|
||||
private def getOrCreateSparkSession(): SparkSession = {
|
||||
SparkSession
|
||||
.builder()
|
||||
.appName("MLflow Spark Autologging Tests")
|
||||
.config("spark.master", "local")
|
||||
.getOrCreate()
|
||||
}
|
||||
|
||||
override def beforeAll(): Unit = {
|
||||
super.beforeAll()
|
||||
// Generate dummy data & write it in various formats (CSV, JSON, parquet)
|
||||
val rows = Seq(
|
||||
Row(8, "bat"),
|
||||
Row(64, "mouse"),
|
||||
Row(-27, "horse")
|
||||
)
|
||||
val schema = List(
|
||||
StructField("number", IntegerType),
|
||||
StructField("word", StringType)
|
||||
)
|
||||
val df = spark.createDataFrame(
|
||||
spark.sparkContext.parallelize(rows),
|
||||
StructType(schema)
|
||||
)
|
||||
tempDir = Files.createTempDirectory(this.getClass.getName)
|
||||
deltaTablePath = Paths.get(tempDir.toString, "delta").toString
|
||||
formatToTablePath = Seq( "csv", "parquet", "json" /*, delta */).map { format =>
|
||||
format -> Paths.get(tempDir.toString, format).toString
|
||||
}.toMap
|
||||
|
||||
formatToTablePath.foreach { case (format, tablePath) =>
|
||||
df.write.option("header", "true").format(format).save(tablePath)
|
||||
}
|
||||
}
|
||||
|
||||
override def afterAll(): Unit = {
|
||||
super.afterAll()
|
||||
def deleteRecursively(file: File): Unit = {
|
||||
if (file.isDirectory) {
|
||||
file.listFiles.foreach(deleteRecursively)
|
||||
}
|
||||
if (file.exists && !file.delete) {
|
||||
throw new RuntimeException(s"Unable to delete ${file.getAbsolutePath}")
|
||||
}
|
||||
}
|
||||
deleteRecursively(tempDir.toFile)
|
||||
}
|
||||
|
||||
override def beforeEach(): Unit = {
|
||||
super.beforeEach()
|
||||
MlflowAutologEventPublisher.init()
|
||||
}
|
||||
|
||||
override def afterEach(): Unit = {
|
||||
MlflowAutologEventPublisher.stop()
|
||||
super.afterEach()
|
||||
}
|
||||
|
||||
private def getFileUri(absolutePath: String): String = {
|
||||
s"${Paths.get("file:", absolutePath).toString}"
|
||||
}
|
||||
|
||||
test("MlflowAutologEventPublisher can be idempotently initialized & stopped within " +
|
||||
"single thread") {
|
||||
// We expect a listener to already be created by calling init() in beforeEach
|
||||
val listeners0 = MlflowSparkAutologgingTestUtils.getListeners(spark)
|
||||
assert(listeners0.length == 1)
|
||||
val listener0 = listeners0.head
|
||||
// Call init() again, verify listener is unchanged
|
||||
MlflowAutologEventPublisher.init()
|
||||
val listeners1 = MlflowSparkAutologgingTestUtils.getListeners(spark)
|
||||
assert(listeners1.length == 1)
|
||||
val listener1 = listeners1.head
|
||||
assert(listener0 == listener1)
|
||||
// Call stop() multiple times
|
||||
MlflowAutologEventPublisher.stop()
|
||||
assert(MlflowSparkAutologgingTestUtils.getListeners(spark).isEmpty)
|
||||
MlflowAutologEventPublisher.stop()
|
||||
assert(MlflowSparkAutologgingTestUtils.getListeners(spark).isEmpty)
|
||||
// Call init() after stop(), verify that we create a new listener
|
||||
MlflowAutologEventPublisher.init()
|
||||
val listeners2 = MlflowSparkAutologgingTestUtils.getListeners(spark)
|
||||
assert(listeners2.length == 1)
|
||||
val listener2 = listeners2.head
|
||||
assert(listener2 != listener1)
|
||||
}
|
||||
|
||||
test("MlflowAutologEventPublisher triggers publishEvent with appropriate arguments " +
|
||||
"when reading datasources corresponding to different formats") {
|
||||
val formatToTestDFs = formatToTablePath.map { case (format, tablePath) =>
|
||||
val baseDf = spark.read.format(format).option("inferSchema", "true")
|
||||
.option("header", "true").load(tablePath)
|
||||
format -> Seq(
|
||||
baseDf,
|
||||
baseDf.filter("number > 0"),
|
||||
baseDf.select("number"),
|
||||
baseDf.limit(2),
|
||||
baseDf.filter("number > 0").select("number").limit(2)
|
||||
)
|
||||
}
|
||||
|
||||
formatToTestDFs.foreach { case (format, dfs) =>
|
||||
dfs.foreach { df =>
|
||||
df.printSchema()
|
||||
MlflowAutologEventPublisher.init()
|
||||
val subscriber = spy(new MockSubscriber())
|
||||
MlflowAutologEventPublisher.register(subscriber)
|
||||
assert(MlflowAutologEventPublisher.subscribers.size == 1)
|
||||
// Read DF
|
||||
df.collect()
|
||||
// Verify events logged
|
||||
Thread.sleep(1000)
|
||||
val tablePath = formatToTablePath(format)
|
||||
val expectedPath = getFileUri(tablePath)
|
||||
verify(subscriber, times(1)).notify(any(), any(), any())
|
||||
verify(subscriber, times(1)).notify(expectedPath, "unknown", format)
|
||||
MlflowAutologEventPublisher.stop()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
test("MlflowAutologEventPublisher triggers publishEvent with appropriate arguments " +
|
||||
"when reading a JOIN of two tables") {
|
||||
val formats = formatToTablePath.keys
|
||||
val leftFormat = formats.head
|
||||
val rightFormat = formats.last
|
||||
val leftPath = formatToTablePath(leftFormat)
|
||||
val rightPath = formatToTablePath(rightFormat)
|
||||
val leftDf = spark.read.format(leftFormat).load(leftPath)
|
||||
val rightDf = spark.read.format(rightFormat).load(rightPath)
|
||||
MlflowAutologEventPublisher.init()
|
||||
val subscriber = spy(new MockSubscriber())
|
||||
MlflowAutologEventPublisher.register(subscriber)
|
||||
leftDf.join(rightDf).collect()
|
||||
// Sleep to let the SparkListener trigger read
|
||||
Thread.sleep(1000)
|
||||
verify(subscriber, times(2)).notify(any(), any(), any())
|
||||
verify(subscriber, times(1)).notify(getFileUri(leftPath), "unknown", leftFormat)
|
||||
verify(subscriber, times(1)).notify(getFileUri(rightPath), "unknown", rightFormat)
|
||||
}
|
||||
|
||||
test("MlflowAutologEventPublisher can publish to working subscribers even when " +
|
||||
"others are broken") {
|
||||
MlflowAutologEventPublisher.stop()
|
||||
val subscriber = spy(new MockSubscriber())
|
||||
// Publish to a broken subscriber, then a working one, and finally another broken one
|
||||
val subscriberSeq = Seq(new BrokenSubscriber(), subscriber, new BrokenSubscriber())
|
||||
object MockPublisher extends MlflowAutologEventPublisherImpl {
|
||||
// Override subscriber iteration logic to yield subscribers in the desired order
|
||||
override def getSubscribers: Seq[(String, MlflowAutologEventSubscriber)] = {
|
||||
subscriberSeq.map(subscriber => (subscriber.replId, subscriber))
|
||||
}
|
||||
}
|
||||
// Disable GC of dead subscribers so that they get published-to
|
||||
MockPublisher.init(gcDeadSubscribersIntervalSec = 10000)
|
||||
val listeners1 = MlflowSparkAutologgingTestUtils.getListeners(spark)
|
||||
assert(listeners1.length == 1)
|
||||
|
||||
val (format, path) = formatToTablePath.head
|
||||
val df = spark.read.format(format).load(path)
|
||||
|
||||
// Register subscribers & collect the DF to trigger a datasource read event
|
||||
subscriberSeq.foreach(MockPublisher.register)
|
||||
df.collect()
|
||||
Thread.sleep(1000)
|
||||
verify(subscriber, times(1)).notify(any(), any(), any())
|
||||
verify(subscriber, times(1)).notify(
|
||||
getFileUri(path), "unknown", format)
|
||||
}
|
||||
|
||||
test("Exceptions while extracting datasource information from Spark query plan " +
|
||||
"do not fail the query") {
|
||||
MlflowAutologEventPublisher.stop()
|
||||
object MockPublisher extends MlflowAutologEventPublisherImpl {
|
||||
// Return a custom listener that throws while processing SparkListenerSQLExecutionEnd events
|
||||
override def getSparkDataSourceListener: SparkDataSourceListener = {
|
||||
new SparkDataSourceListener {
|
||||
override def onSQLExecutionEnd(event: SparkListenerSQLExecutionEnd): Unit = {
|
||||
throw new NoSuchMethodException("Mock failure while extracting datasource info from " +
|
||||
"query plan!")
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
MockPublisher.init()
|
||||
val (format, path) = formatToTablePath.head
|
||||
val df = spark.read.format(format).load(path)
|
||||
val subscriber = new MockSubscriber()
|
||||
MockPublisher.register(subscriber)
|
||||
df.collect()
|
||||
}
|
||||
|
||||
test("ReplAwareDatasourceAttributeExtractor handles missing Databricks classes gracefully") {
|
||||
import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
|
||||
import org.mockito.ArgumentCaptor
|
||||
import scala.util.control.NonFatal
|
||||
|
||||
MlflowAutologEventPublisher.stop()
|
||||
|
||||
var deltaDetectionAttempted = false
|
||||
var exceptionCaught = false
|
||||
|
||||
object TrackingDatasourceAttributeExtractor extends DatasourceAttributeExtractorBase {
|
||||
override protected def maybeGetDeltaTableInfo(leafNode: LogicalPlan): Option[SparkTableInfo] = {
|
||||
deltaDetectionAttempted = true
|
||||
try {
|
||||
ReflectionUtils.getScalaObjectByName(
|
||||
"com.databricks.sql.transaction.tahoe.DeltaTable")
|
||||
throw new AssertionError(
|
||||
"Databricks Delta class unexpectedly found - this test should run without it")
|
||||
} catch {
|
||||
case NonFatal(_) =>
|
||||
exceptionCaught = true
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
class ReplAwareListenerWithTrackingExtractor(
|
||||
publisher: MlflowAutologEventPublisherImpl = MlflowAutologEventPublisher)
|
||||
extends ReplAwareSparkDataSourceListener(publisher) {
|
||||
override protected def getDatasourceAttributeExtractor: DatasourceAttributeExtractorBase = {
|
||||
TrackingDatasourceAttributeExtractor
|
||||
}
|
||||
}
|
||||
|
||||
object MockPublisher extends MlflowAutologEventPublisherImpl {
|
||||
override def getSparkDataSourceListener: SparkDataSourceListener = {
|
||||
new ReplAwareListenerWithTrackingExtractor(this)
|
||||
}
|
||||
}
|
||||
|
||||
MockPublisher.init()
|
||||
val (format, path) = formatToTablePath.head
|
||||
val df = spark.read.format(format).load(path)
|
||||
val subscriber = spy(new MockSubscriber())
|
||||
MockPublisher.register(subscriber)
|
||||
|
||||
val sc = spark.sparkContext
|
||||
sc.setLocalProperty("spark.databricks.replId", subscriber.replId)
|
||||
|
||||
df.collect()
|
||||
|
||||
Thread.sleep(1000)
|
||||
assert(deltaDetectionAttempted, "Delta detection should have been attempted")
|
||||
assert(exceptionCaught,
|
||||
"Exception should have been caught when loading missing Databricks class")
|
||||
|
||||
val formatCaptor = ArgumentCaptor.forClass(classOf[String])
|
||||
verify(subscriber, times(1)).notify(any(), any(), formatCaptor.capture())
|
||||
assert(formatCaptor.getValue != "delta",
|
||||
"Format should not be 'delta' since Databricks classes are unavailable")
|
||||
}
|
||||
|
||||
test("MlflowAutologEventPublisher correctly unregisters broken subscribers") {
|
||||
MlflowAutologEventPublisher.register(new BrokenSubscriber())
|
||||
Thread.sleep(2000)
|
||||
assert(MlflowAutologEventPublisher.subscribers.isEmpty)
|
||||
}
|
||||
|
||||
test("Subscriber registration fails if init() not called") {
|
||||
MlflowAutologEventPublisher.stop()
|
||||
intercept[RuntimeException] {
|
||||
MlflowAutologEventPublisher.register(new MockSubscriber())
|
||||
}
|
||||
}
|
||||
|
||||
test("Initializing MlflowAutologEventPublisher fails if SparkSession doesn't exixt") {
|
||||
MlflowAutologEventPublisher.stop()
|
||||
spark.stop()
|
||||
try {
|
||||
intercept[RuntimeException] {
|
||||
MlflowAutologEventPublisher.init()
|
||||
}
|
||||
} finally {
|
||||
spark = getOrCreateSparkSession()
|
||||
}
|
||||
}
|
||||
|
||||
test("Delegates to repl-ID-aware listener if REPL ID property is set in SparkContext") {
|
||||
// Verify instance created by init() in beforeEach is not REPL-ID-aware
|
||||
assert(MlflowAutologEventPublisher.sparkQueryListener.isInstanceOf[SparkDataSourceListener])
|
||||
assert(!MlflowAutologEventPublisher.sparkQueryListener.isInstanceOf[ReplAwareSparkDataSourceListener])
|
||||
// Call stop, update SparkContext to contain repl ID property, call init(), verify instance is
|
||||
// REPL-ID-aware
|
||||
MlflowAutologEventPublisher.stop()
|
||||
assert(MlflowSparkAutologgingTestUtils.getListeners(spark).isEmpty)
|
||||
val sc = spark.sparkContext
|
||||
sc.setLocalProperty("spark.databricks.replId", "myCoolReplId")
|
||||
MlflowAutologEventPublisher.init()
|
||||
assert(MlflowAutologEventPublisher.sparkQueryListener.isInstanceOf[ReplAwareSparkDataSourceListener])
|
||||
sc.setLocalProperty("spark.databricks.replId", null)
|
||||
MlflowAutologEventPublisher.stop()
|
||||
MlflowAutologEventPublisher.init()
|
||||
assert(MlflowAutologEventPublisher.sparkQueryListener.isInstanceOf[SparkDataSourceListener])
|
||||
}
|
||||
|
||||
test("Delegates to repl-ID-aware listener if Databricks cluster ID is set in Spark Conf") {
|
||||
// Verify instance created by init() in beforeEach is not REPL-ID-aware
|
||||
assert(MlflowAutologEventPublisher.sparkQueryListener.isInstanceOf[SparkDataSourceListener])
|
||||
assert(!MlflowAutologEventPublisher.sparkQueryListener.isInstanceOf[ReplAwareSparkDataSourceListener])
|
||||
MlflowAutologEventPublisher.stop()
|
||||
|
||||
spark.conf.set("spark.databricks.clusterUsageTags.clusterId", "myCoolClusterId")
|
||||
MlflowAutologEventPublisher.init()
|
||||
assert(MlflowAutologEventPublisher.sparkQueryListener.isInstanceOf[ReplAwareSparkDataSourceListener])
|
||||
}
|
||||
|
||||
test("repl-ID-aware listener publishes events with expected REPL IDs") {
|
||||
MlflowAutologEventPublisher.stop()
|
||||
|
||||
// Create a ReplAwareSparkDataSourceListener that uses a DatasourceAttributeExtractor instead
|
||||
// of a ReplAwareDatasourceAttributeExtractor for testing, since
|
||||
// ReplAwareDatasourceAttributeExtractor requires Databricks-specific packages that are not
|
||||
// available in OSS test environments
|
||||
class ReplAwareSparkDataSourceListenerWithDefaultDatasourceAttributeExtractor(
|
||||
publisher: MlflowAutologEventPublisherImpl = MlflowAutologEventPublisher)
|
||||
extends ReplAwareSparkDataSourceListener(publisher) {
|
||||
override protected def getDatasourceAttributeExtractor: DatasourceAttributeExtractorBase = {
|
||||
DatasourceAttributeExtractor
|
||||
}
|
||||
}
|
||||
|
||||
// Create and initialize a publisher that uses the ReplAwareSparkDataSourceListener containing
|
||||
// the DatasourceAttributeExtractor defined above
|
||||
object MockReplAwarePublisher extends MlflowAutologEventPublisherImpl {
|
||||
override def getSparkDataSourceListener: SparkDataSourceListener = {
|
||||
new ReplAwareSparkDataSourceListenerWithDefaultDatasourceAttributeExtractor(this)
|
||||
}
|
||||
}
|
||||
MockReplAwarePublisher.init()
|
||||
// Register several subcribers with different REPL IDs
|
||||
val subscriber1 = spy(new MockSubscriber())
|
||||
val subscriber2 = spy(new MockSubscriber())
|
||||
val subscriber3 = spy(new MockSubscriber())
|
||||
MockReplAwarePublisher.register(subscriber1)
|
||||
MockReplAwarePublisher.register(subscriber2)
|
||||
MockReplAwarePublisher.register(subscriber3)
|
||||
|
||||
val sc = spark.sparkContext
|
||||
val formatToTablePathList = formatToTablePath.toList
|
||||
|
||||
// Read a collection of Spark DataFrames from different sources with different REPL ID
|
||||
// context for each read
|
||||
|
||||
// Because `spark.databricks.replId` is null, we expect that none of the subscribers will
|
||||
// be notified when `path1` is read via `df1`
|
||||
sc.setLocalProperty("spark.databricks.replId", null)
|
||||
val (format1, path1) = formatToTablePathList.head
|
||||
val df1 = spark.read.format(format1).load(path1)
|
||||
df1.collect()
|
||||
|
||||
// Because `spark.databricks.replId` is set to `subscriber1.replId`, we expect that only
|
||||
// `subscriber1` will be notified when `path2` is read via `df2`
|
||||
sc.setLocalProperty("spark.databricks.replId", subscriber1.replId)
|
||||
val (format2, path2) = formatToTablePathList(1)
|
||||
val df2 = spark.read.format(format2).load(path2)
|
||||
df2.collect()
|
||||
|
||||
// Because `spark.databricks.replId` is set to `subscriber2.replId`, we expect that only
|
||||
// `subscriber2` will be notified when `path3` is read via `df3`
|
||||
sc.setLocalProperty("spark.databricks.replId", subscriber2.replId)
|
||||
val (format3, path3) = formatToTablePathList(2)
|
||||
val df3 = spark.read.format(format3).load(path3)
|
||||
df3.collect()
|
||||
|
||||
// Because `spark.databricks.replId` is set to `subscriber3.replId`, we expect that only
|
||||
// `subscriber3` will be notified when `path1`, `path2`, and `path3` are read via `df4`
|
||||
sc.setLocalProperty("spark.databricks.replId", subscriber3.replId)
|
||||
val df4 = df1.union(df2).union(df3)
|
||||
df4.collect()
|
||||
|
||||
// Sleep to give time for the execution to complete
|
||||
Thread.sleep(1000)
|
||||
|
||||
// Verify that the only time subscriber1 was notified of a datasource read was when
|
||||
// `path2` was read via `df2` with `spark.databricks.replId` set to `subscriber1.replId`
|
||||
verify(subscriber1, times(1)).notify(any(), any(), any())
|
||||
verify(subscriber1, times(1)).notify(meq(s"file:$path2"), any(), meq(format2))
|
||||
|
||||
// Verify that the only time subscriber2 was notified of a datasource read was when
|
||||
// `path3` was read via `df3` with `spark.databricks.replId` set to `subscriber2.replId`
|
||||
verify(subscriber2, times(1)).notify(any(), any(), any())
|
||||
verify(subscriber2, times(1)).notify(meq(s"file:$path3"), any(), meq(format3))
|
||||
|
||||
// Verify that subscriber3 was notified of three datasource reads - one for each of
|
||||
// `path1`, `path2`, and `path3` - via `df4` with `spark.databricks.replId` set to
|
||||
// `subscriber3.replId`
|
||||
verify(subscriber3, times(3)).notify(any(), any(), any())
|
||||
verify(subscriber3, times(1)).notify(meq(s"file:$path1"), any(), meq(format1))
|
||||
verify(subscriber3, times(1)).notify(meq(s"file:$path2"), any(), meq(format2))
|
||||
verify(subscriber3, times(1)).notify(meq(s"file:$path3"), any(), meq(format3))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,134 @@
|
||||
<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
|
||||
<modelVersion>4.0.0</modelVersion>
|
||||
<artifactId>mlflow-spark_2.12</artifactId>
|
||||
<version>3.14.1-SNAPSHOT</version>
|
||||
<name>${project.artifactId}</name>
|
||||
<properties>
|
||||
<maven.compiler.source>1.8</maven.compiler.source>
|
||||
<maven.compiler.target>1.8</maven.compiler.target>
|
||||
<encoding>UTF-8</encoding>
|
||||
<scala.version>2.12.18</scala.version>
|
||||
<scala.compat.version>2.12</scala.compat.version>
|
||||
<spec2.version>4.2.0</spec2.version>
|
||||
</properties>
|
||||
|
||||
<parent>
|
||||
<groupId>org.mlflow</groupId>
|
||||
<artifactId>mlflow-parent</artifactId>
|
||||
<version>3.14.1-SNAPSHOT</version>
|
||||
<relativePath>../pom.xml</relativePath>
|
||||
</parent>
|
||||
|
||||
<dependencies>
|
||||
<dependency>
|
||||
<groupId>org.scala-lang</groupId>
|
||||
<artifactId>scala-library</artifactId>
|
||||
<version>2.12.18</version>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.apache.spark</groupId>
|
||||
<artifactId>spark-sql_2.12</artifactId>
|
||||
<version>3.5.0</version>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.apache.spark</groupId>
|
||||
<artifactId>spark-core_2.12</artifactId>
|
||||
<version>3.5.0</version>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.apache.spark</groupId>
|
||||
<artifactId>spark-catalyst_2.12</artifactId>
|
||||
<version>3.5.0</version>
|
||||
</dependency>
|
||||
|
||||
<!-- Test -->
|
||||
<dependency>
|
||||
<groupId>junit</groupId>
|
||||
<artifactId>junit</artifactId>
|
||||
<version>4.12</version>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.mockito</groupId>
|
||||
<artifactId>mockito-all</artifactId>
|
||||
<version>1.10.19</version>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.scalatest</groupId>
|
||||
<artifactId>scalatest_2.12</artifactId>
|
||||
<version>3.2.17</version>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
|
||||
</dependencies>
|
||||
|
||||
<build>
|
||||
<sourceDirectory>../spark/src/main/scala</sourceDirectory>
|
||||
<testSourceDirectory>../spark/src/test/scala</testSourceDirectory>
|
||||
<plugins>
|
||||
<plugin>
|
||||
<groupId>net.alchim31.maven</groupId>
|
||||
<artifactId>scala-maven-plugin</artifactId>
|
||||
<version>4.8.1</version>
|
||||
<executions>
|
||||
<execution>
|
||||
<goals>
|
||||
<goal>compile</goal>
|
||||
<goal>testCompile</goal>
|
||||
</goals>
|
||||
<configuration>
|
||||
<args>
|
||||
<arg>-dependencyfile</arg>
|
||||
<arg>${project.build.directory}/.scala_dependencies</arg>
|
||||
</args>
|
||||
</configuration>
|
||||
</execution>
|
||||
<execution>
|
||||
<id>attach-javadocs</id>
|
||||
<goals>
|
||||
<goal>doc-jar</goal>
|
||||
</goals>
|
||||
</execution>
|
||||
</executions>
|
||||
</plugin>
|
||||
<plugin>
|
||||
<groupId>org.apache.maven.plugins</groupId>
|
||||
<artifactId>maven-surefire-plugin</artifactId>
|
||||
<version>2.7</version>
|
||||
<configuration>
|
||||
<skipTests>true</skipTests>
|
||||
</configuration>
|
||||
</plugin>
|
||||
<plugin>
|
||||
<groupId>org.scalatest</groupId>
|
||||
<artifactId>scalatest-maven-plugin</artifactId>
|
||||
<version>1.0</version>
|
||||
<configuration>
|
||||
<reportsDirectory>${project.build.directory}/surefire-reports</reportsDirectory>
|
||||
<junitxml>.</junitxml>
|
||||
<filereports>WDF TestSuite.txt</filereports>
|
||||
</configuration>
|
||||
<executions>
|
||||
<execution>
|
||||
<id>test</id>
|
||||
<goals>
|
||||
<goal>test</goal>
|
||||
</goals>
|
||||
</execution>
|
||||
</executions>
|
||||
</plugin>
|
||||
<plugin>
|
||||
<groupId>org.apache.maven.plugins</groupId>
|
||||
<artifactId>maven-source-plugin</artifactId>
|
||||
</plugin>
|
||||
<plugin>
|
||||
<groupId>org.apache.maven.plugins</groupId>
|
||||
<artifactId>maven-javadoc-plugin</artifactId>
|
||||
<configuration>
|
||||
<sourcepath>${project.basedir}/src/main/scala</sourcepath>
|
||||
</configuration>
|
||||
</plugin>
|
||||
</plugins>
|
||||
</build>
|
||||
</project>
|
||||
@@ -0,0 +1,134 @@
|
||||
<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
|
||||
<modelVersion>4.0.0</modelVersion>
|
||||
<artifactId>mlflow-spark_2.13</artifactId>
|
||||
<version>3.14.1-SNAPSHOT</version>
|
||||
<name>${project.artifactId}</name>
|
||||
<properties>
|
||||
<maven.compiler.source>1.8</maven.compiler.source>
|
||||
<maven.compiler.target>1.8</maven.compiler.target>
|
||||
<encoding>UTF-8</encoding>
|
||||
<scala.version>2.13.10</scala.version>
|
||||
<scala.compat.version>2.13</scala.compat.version>
|
||||
<spec2.version>4.2.0</spec2.version>
|
||||
</properties>
|
||||
|
||||
<parent>
|
||||
<groupId>org.mlflow</groupId>
|
||||
<artifactId>mlflow-parent</artifactId>
|
||||
<version>3.14.1-SNAPSHOT</version>
|
||||
<relativePath>../pom.xml</relativePath>
|
||||
</parent>
|
||||
|
||||
<dependencies>
|
||||
<dependency>
|
||||
<groupId>org.scala-lang</groupId>
|
||||
<artifactId>scala-library</artifactId>
|
||||
<version>2.13.10</version>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.apache.spark</groupId>
|
||||
<artifactId>spark-sql_2.13</artifactId>
|
||||
<version>3.5.0</version>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.apache.spark</groupId>
|
||||
<artifactId>spark-core_2.13</artifactId>
|
||||
<version>3.5.0</version>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.apache.spark</groupId>
|
||||
<artifactId>spark-catalyst_2.13</artifactId>
|
||||
<version>3.5.0</version>
|
||||
</dependency>
|
||||
|
||||
<!-- Test -->
|
||||
<dependency>
|
||||
<groupId>junit</groupId>
|
||||
<artifactId>junit</artifactId>
|
||||
<version>4.12</version>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.mockito</groupId>
|
||||
<artifactId>mockito-all</artifactId>
|
||||
<version>1.10.19</version>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.scalatest</groupId>
|
||||
<artifactId>scalatest_2.13</artifactId>
|
||||
<version>3.2.17</version>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
|
||||
</dependencies>
|
||||
|
||||
<build>
|
||||
<sourceDirectory>../spark/src/main/scala</sourceDirectory>
|
||||
<testSourceDirectory>../spark/src/test/scala</testSourceDirectory>
|
||||
<plugins>
|
||||
<plugin>
|
||||
<groupId>net.alchim31.maven</groupId>
|
||||
<artifactId>scala-maven-plugin</artifactId>
|
||||
<version>4.8.1</version>
|
||||
<executions>
|
||||
<execution>
|
||||
<goals>
|
||||
<goal>compile</goal>
|
||||
<goal>testCompile</goal>
|
||||
</goals>
|
||||
<configuration>
|
||||
<args>
|
||||
<arg>-dependencyfile</arg>
|
||||
<arg>${project.build.directory}/.scala_dependencies</arg>
|
||||
</args>
|
||||
</configuration>
|
||||
</execution>
|
||||
<execution>
|
||||
<id>attach-javadocs</id>
|
||||
<goals>
|
||||
<goal>doc-jar</goal>
|
||||
</goals>
|
||||
</execution>
|
||||
</executions>
|
||||
</plugin>
|
||||
<plugin>
|
||||
<groupId>org.apache.maven.plugins</groupId>
|
||||
<artifactId>maven-surefire-plugin</artifactId>
|
||||
<version>2.7</version>
|
||||
<configuration>
|
||||
<skipTests>true</skipTests>
|
||||
</configuration>
|
||||
</plugin>
|
||||
<plugin>
|
||||
<groupId>org.scalatest</groupId>
|
||||
<artifactId>scalatest-maven-plugin</artifactId>
|
||||
<version>1.0</version>
|
||||
<configuration>
|
||||
<reportsDirectory>${project.build.directory}/surefire-reports</reportsDirectory>
|
||||
<junitxml>.</junitxml>
|
||||
<filereports>WDF TestSuite.txt</filereports>
|
||||
</configuration>
|
||||
<executions>
|
||||
<execution>
|
||||
<id>test</id>
|
||||
<goals>
|
||||
<goal>test</goal>
|
||||
</goals>
|
||||
</execution>
|
||||
</executions>
|
||||
</plugin>
|
||||
<plugin>
|
||||
<groupId>org.apache.maven.plugins</groupId>
|
||||
<artifactId>maven-source-plugin</artifactId>
|
||||
</plugin>
|
||||
<plugin>
|
||||
<groupId>org.apache.maven.plugins</groupId>
|
||||
<artifactId>maven-javadoc-plugin</artifactId>
|
||||
<configuration>
|
||||
<sourcepath>${project.basedir}/src/main/scala</sourcepath>
|
||||
</configuration>
|
||||
</plugin>
|
||||
</plugins>
|
||||
</build>
|
||||
</project>
|
||||
Reference in New Issue
Block a user